# SPDX-License-Identifier: Apache-2.0 """Tests for the GLM-5.2 (glm_moe_dsa) native MTP patch.""" import sys import mlx.core as mx import mlx.utils as mu import pytest from omlx.patches.glm_moe_dsa import apply_glm_moe_dsa_patch from omlx.patches.mlx_lm_mtp import apply_mlx_lm_mtp_patch, set_mtp_active @pytest.fixture(scope="module") def glm(): apply_glm_moe_dsa_patch() apply_mlx_lm_mtp_patch() return sys.modules["mlx_lm.models.glm_moe_dsa"] @pytest.fixture() def mtp_active(): set_mtp_active(True) yield set_mtp_active(False) TINY_CFG = dict( model_type="glm_moe_dsa", vocab_size=128, hidden_size=64, index_head_dim=32, index_n_heads=4, index_topk=16, intermediate_size=96, moe_intermediate_size=32, num_hidden_layers=2, num_attention_heads=4, num_key_value_heads=4, n_shared_experts=1, n_routed_experts=4, routed_scaling_factor=1.0, kv_lora_rank=32, q_lora_rank=48, qk_rope_head_dim=16, v_head_dim=32, qk_nope_head_dim=24, topk_method="noaux_tc", scoring_func="sigmoid", norm_topk_prob=True, n_group=1, topk_group=1, num_experts_per_tok=2, moe_layer_freq=1, first_k_dense_replace=1, max_position_embeddings=512, rms_norm_eps=1e-5, rope_parameters={"rope_theta": 10000.0, "rope_type": "default"}, attention_bias=False, index_topk_freq=4, index_skip_topk_offset=3, indexer_types=["full", "shared"], num_nextn_predict_layers=1, ) def _raw_hf_weights(glm, model): """Rebuild a raw-HF-layout weights dict from a built model's params. Inverts the sanitize transforms for the MTP layer (switch stacking, gate_up fusion, embed_q/unembed_out) so sanitize can be exercised on checkpoint-shaped input. """ cfg = TINY_CFG flat = dict(mu.tree_flatten(model.parameters())) weights = {} for k, v in flat.items(): if k.startswith("mtp.0."): rest = k[len("mtp.0."):] if rest.startswith("block."): rk = "model.layers.2." + rest[len("block."):] elif rest == "norm.weight": rk = "model.layers.2.shared_head.norm.weight" else: rk = "model.layers.2." + rest weights[rk] = v else: weights[k] = v raw = {} for k, v in weights.items(): if ".mlp.switch_mlp.gate_up_proj.weight" in k: base = k.split(".mlp.switch_mlp.")[0] gate, up = mx.split(v, 2, axis=1) for e in range(v.shape[0]): raw[f"{base}.mlp.experts.{e}.gate_proj.weight"] = gate[e] raw[f"{base}.mlp.experts.{e}.up_proj.weight"] = up[e] elif ".mlp.switch_mlp.down_proj.weight" in k: base = k.split(".mlp.switch_mlp.")[0] for e in range(v.shape[0]): raw[f"{base}.mlp.experts.{e}.down_proj.weight"] = v[e] elif ".self_attn.embed_q.weight" in k: continue # regenerated from the fabricated kv_b_proj below elif ".self_attn.unembed_out.weight" in k: base = k.split(".self_attn.")[0] nh = cfg["num_attention_heads"] hd = cfg["qk_nope_head_dim"] + cfg["v_head_dim"] raw[f"{base}.self_attn.kv_b_proj.weight"] = mx.random.normal( (nh * hd, cfg["kv_lora_rank"]) ) else: raw[k] = v return raw class TestModelArgs: def test_nextn_count_and_indexer_extension(self, glm): args = glm.ModelArgs.from_dict(TINY_CFG) assert args.num_nextn_predict_layers == 1 # freq=4/offset=3: layer 2 -> max(0,0)%4==0 -> "full" assert args.indexer_types == ["full", "shared", "full"] def test_no_nextn_is_untouched(self, glm): cfg = dict(TINY_CFG, num_nextn_predict_layers=0) args = glm.ModelArgs.from_dict(cfg) assert args.num_nextn_predict_layers == 0 assert args.indexer_types == ["full", "shared"] class TestModelInit: def test_mtp_attached_when_active(self, glm, mtp_active): args = glm.ModelArgs.from_dict(TINY_CFG) model = glm.Model(args) assert hasattr(model, "mtp") and len(model.mtp) == 1 assert model._omlx_mtp_decode_enabled assert model._omlx_mtp_chain assert model._omlx_mtp_head_clone is False def test_mtp_skipped_when_inactive(self, glm): set_mtp_active(False) args = glm.ModelArgs.from_dict(TINY_CFG) model = glm.Model(args) assert not hasattr(model, "mtp") assert model._omlx_mtp_decode_enabled is False class TestSanitize: def test_raw_hf_remap_and_strict_load(self, glm, mtp_active): mx.random.seed(0) args = glm.ModelArgs.from_dict(TINY_CFG) model = glm.Model(args) raw = _raw_hf_weights(glm, model) out = model.sanitize(raw) assert not any(".layers.2." in k for k in out) for expected in ( "mtp.0.eh_proj.weight", "mtp.0.enorm.weight", "mtp.0.hnorm.weight", "mtp.0.norm.weight", "mtp.0.block.mlp.switch_mlp.gate_up_proj.weight", "mtp.0.block.self_attn.embed_q.weight", "mtp.0.block.self_attn.indexer.wk.weight", ): assert expected in out, expected model.load_weights(list(out.items()), strict=True) def test_layer_count_restored_after_sanitize(self, glm, mtp_active): args = glm.ModelArgs.from_dict(TINY_CFG) model = glm.Model(args) raw = _raw_hf_weights(glm, model) model.sanitize(raw) assert model.args.num_hidden_layers == TINY_CFG["num_hidden_layers"] def test_presanitized_passthrough(self, glm, mtp_active): """oQ-style checkpoints (already mtp.*) survive a second sanitize.""" mx.random.seed(0) args = glm.ModelArgs.from_dict(TINY_CFG) model = glm.Model(args) once = model.sanitize(_raw_hf_weights(glm, model)) twice = model.sanitize(dict(once)) assert sorted(twice) == sorted(once) model.load_weights(list(twice.items()), strict=True) def test_mtp_off_drops_all_mtp_keys(self, glm, mtp_active): mx.random.seed(0) args = glm.ModelArgs.from_dict(TINY_CFG) model = glm.Model(args) sanitized = model.sanitize(_raw_hf_weights(glm, model)) set_mtp_active(False) model_off = glm.Model(args) out = model_off.sanitize(dict(sanitized)) assert not any(k.startswith("mtp.") for k in out) model_off.load_weights(list(out.items()), strict=True) def test_missing_head_weights_degrades_gracefully(self, glm, mtp_active): mx.random.seed(0) args = glm.ModelArgs.from_dict(TINY_CFG) model = glm.Model(args) sanitized = model.sanitize(_raw_hf_weights(glm, model)) stripped = {k: v for k, v in sanitized.items() if not k.startswith("mtp.")} model2 = glm.Model(args) out = model2.sanitize(stripped) assert not hasattr(model2, "mtp") assert model2._omlx_mtp_decode_enabled is False model2.load_weights(list(out.items()), strict=True) class TestIndexerFusion: def test_mtp_indexer_fused_alongside_backbone(self, glm, mtp_active): """MTP indexer fusion must happen before the stock sanitize: its backbone fusion pass drops every unfused ``.indexer.wk`` / ``.weights_proj`` key by substring, MTP keys included.""" cfg = dict(TINY_CFG) q8 = {"bits": 8, "group_size": 64, "mode": "affine"} cfg["quantization"] = { "group_size": 64, "bits": 4, "mode": "affine", "model.layers.0.self_attn.indexer.wk": dict(q8), "model.layers.0.self_attn.indexer.weights_proj": dict(q8), } args = glm.ModelArgs.from_dict(cfg) model = glm.Model(args) assert model.mtp[0].block.self_attn.indexer.wk_weights_proj is not None h = TINY_CFG["hidden_size"] hd = TINY_CFG["index_head_dim"] nh = TINY_CFG["index_n_heads"] weights = {} for prefix in ( "model.layers.0.self_attn.indexer", "mtp.0.block.self_attn.indexer", ): for suffix in ("weight", "scales", "biases"): weights[f"{prefix}.wk.{suffix}"] = mx.zeros((hd, 4)) weights[f"{prefix}.weights_proj.{suffix}"] = mx.zeros((nh, 4)) out = model.sanitize(weights) for prefix in ( "model.layers.0.self_attn.indexer", "mtp.0.block.self_attn.indexer", ): assert f"{prefix}.wk_weights_proj.weight" in out, prefix assert f"{prefix}.wk.weight" not in out assert f"{prefix}.weights_proj.weight" not in out assert out["mtp.0.block.self_attn.indexer.wk_weights_proj.weight"].shape == ( hd + nh, 4, ) class TestForward: @pytest.fixture() def loaded(self, glm, mtp_active): mx.random.seed(0) args = glm.ModelArgs.from_dict(TINY_CFG) model = glm.Model(args) out = model.sanitize(_raw_hf_weights(glm, model)) model.load_weights(list(out.items()), strict=True) mx.eval(model.parameters()) return model def test_return_hidden_and_mtp_cycle(self, loaded): model = loaded cache = model.make_cache() toks = mx.array([[1, 2, 3, 4]]) logits, hidden = model(toks, cache=cache, return_hidden=True) mx.eval(logits, hidden) assert logits.shape == (1, 4, TINY_CFG["vocab_size"]) assert hidden.shape == (1, 4, TINY_CFG["hidden_size"]) # hidden is pre-norm: normed hidden feeds the head (post-norm contract) post = model.model.norm(hidden) mtp_cache = model.make_mtp_cache() assert isinstance(mtp_cache, list) and len(mtp_cache) == 2 lg, hh = model.mtp_forward( post, toks, mtp_cache, return_hidden=True, logits_keep=1 ) mx.eval(lg, hh) assert lg.shape == (1, 1, TINY_CFG["vocab_size"]) assert hh.shape == (1, 4, TINY_CFG["hidden_size"]) assert mtp_cache[0].offset == 4 and mtp_cache[1].offset == 4 # chained draft step + rollback trim lg2, _ = model.mtp_forward( hh[:, -1:], mx.array([[7]]), mtp_cache, return_hidden=True ) mx.eval(lg2) assert mtp_cache[0].offset == 5 from omlx.patches.mlx_lm_mtp.batch_generator import _mtp_head_trim_to _mtp_head_trim_to(mtp_cache, 4) assert mtp_cache[0].offset == 4 and mtp_cache[1].offset == 4 def test_partial_rollback_trims_verify_window(self, loaded): model = loaded cache = model.make_cache() logits, _ = model(mx.array([[1, 2, 3]]), cache=cache, return_hidden=True) mx.eval(logits, *(c[0].keys for c in cache)) base = cache[0][0].offset # verify window: num_drafts + 1 rows, accept 1 of 3 drafts logits, _ = model( mx.array([[4, 5, 6, 7]]), cache=cache, return_hidden=True ) mx.eval(logits, *(c[0].keys for c in cache)) assert cache[0][0].offset == base + 4 assert model.mtp_partial_rollback(cache, 1, 3) for c in cache: for sub in c.caches: # latent KV (+ indexer KV on full layers) assert sub.offset == base + 2 # next_main + 1 accepted draft def test_n_confirmed_accepted(self, loaded): cache = loaded.make_cache() logits, _ = loaded( mx.array([[1, 2]]), cache=cache, return_hidden=True, n_confirmed=1 ) mx.eval(logits) assert logits.shape == (1, 2, TINY_CFG["vocab_size"]) class TestSmallLRouting: def test_absorbed_matches_materialized(self, glm, mtp_active): """The widened L<=8 absorbed path equals the legacy materialize path.""" import omlx.patches.glm_moe_dsa.glm_moe_dsa_model as gm from mlx_lm.models.base import create_attention_mask from mlx_lm.models.cache import KVCache mx.random.seed(3) args = glm.ModelArgs.from_dict(TINY_CFG) attn = glm.GlmMoeDsaAttention(args, 0) mx.eval(attn.parameters()) def run(L, max_l): mx.random.seed(11) cache = [KVCache(), KVCache()] x_pre = mx.random.normal((1, 12, TINY_CFG["hidden_size"])) mask = create_attention_mask(x_pre, cache[0], return_array=True) out, _ = attn(x_pre, mask, cache, None) mx.eval(out) x = mx.random.normal((1, L, TINY_CFG["hidden_size"])) mask = create_attention_mask(x, cache[0], return_array=True) saved = gm._ABSORBED_DECODE_MAX_L gm._ABSORBED_DECODE_MAX_L = max_l try: out, _ = attn(x, mask, cache, None) mx.eval(out) finally: gm._ABSORBED_DECODE_MAX_L = saved return out for L in (2, 3, 4, 8): legacy = run(L, 1) absorbed = run(L, 8) diff = float(mx.abs(legacy - absorbed).max()) assert diff < 2e-5, f"L={L}: {diff}" def test_topk_gather_matches_masked_reference(self, glm, mtp_active): """With the DSA indexer active (K > index_topk), the decode-shape per-row gather path must equal the legacy masked full-K path.""" import omlx.patches.glm_moe_dsa.glm_moe_dsa_model as gm from mlx_lm.models.base import create_attention_mask from mlx_lm.models.cache import KVCache mx.random.seed(5) args = glm.ModelArgs.from_dict(TINY_CFG) attn = glm.GlmMoeDsaAttention(args, 0) mx.eval(attn.parameters()) def run(L, max_l): mx.random.seed(17) cache = [KVCache(), KVCache()] # Prefill past index_topk (16) so the indexer emits topk state. x_pre = mx.random.normal((1, 24, TINY_CFG["hidden_size"])) mask = create_attention_mask(x_pre, cache[0], return_array=True) out, _ = attn(x_pre, mask, cache, None) mx.eval(out) x = mx.random.normal((1, L, TINY_CFG["hidden_size"])) mask = create_attention_mask(x, cache[0], return_array=True) saved = gm._ABSORBED_DECODE_MAX_L gm._ABSORBED_DECODE_MAX_L = max_l try: out, state = attn(x, mask, cache, None) mx.eval(out) finally: gm._ABSORBED_DECODE_MAX_L = saved return out, state for L in (2, 3, 4): legacy, legacy_state = run(L, 1) # masked materialize fallback gathered, state = run(L, 8) # per-row topk gather idx, prefix = gm._parse_topk_state(state) assert idx is not None and idx.shape[2] == L and prefix == 0 diff = float(mx.abs(legacy - gathered).max()) assert diff < 2e-5, f"L={L}: {diff}"