From 3cb18b51892a7800678923116f5b4920a7666208 Mon Sep 17 00:00:00 2001 From: Pedro Cuenca Date: Tue, 16 Jun 2026 22:29:44 +0200 Subject: [PATCH] glm_moe_dsa: DSA cross-layer indexer sharing (GLM-5.2) GLM-5.2 drives the DSA indexer with a per-layer schedule (config.indexer_types): full layers run their own indexer; shared layers reuse the previous full layer's top-k selection and carry no indexer weights. Implement the schedule and thread prev_topk_indices through the decoder layers; shared layers also skip the indexer KV cache (an unused one crashes on .state during generation). Add a model test covering the full/shared schedule, the sparse path, and cache serialization. Defaults (no schedule) reduce to plain deepseek_v32. Co-Authored-By: Claude Opus 4.8 --- mlx_lm/models/glm_moe_dsa.py | 209 ++++++++++++++++++++++++++++++++++- tests/test_models.py | 60 ++++++++++ 2 files changed, 267 insertions(+), 2 deletions(-) diff --git a/mlx_lm/models/glm_moe_dsa.py b/mlx_lm/models/glm_moe_dsa.py index 14e96e365..639c10274 100644 --- a/mlx_lm/models/glm_moe_dsa.py +++ b/mlx_lm/models/glm_moe_dsa.py @@ -1,9 +1,17 @@ # Copyright © 2025 Apple Inc. from dataclasses import dataclass -from typing import Any, Dict, Optional +from typing import Any, Dict, List, Optional -from .base import BaseModelArgs +import mlx.core as mx + +from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention +from .cache import CacheList, KVCache +from .deepseek_v32 import ( + DeepseekV32Attention, + DeepseekV32DecoderLayer, + DeepseekV32Model, +) from .deepseek_v32 import Model as DSV32Model @@ -42,12 +50,209 @@ class ModelArgs(BaseModelArgs): attention_bias: bool rope_scaling: Dict = None rope_theta: Optional[float] = None + indexer_types: Optional[List[str]] = None + index_topk_pattern: Optional[Any] = None + index_topk_freq: int = 1 + index_skip_topk_offset: int = 2 def __post_init__(self): self.rope_scaling = self.rope_parameters self.rope_theta = self.rope_parameters["rope_theta"] + if self.indexer_types is None: + if self.index_topk_pattern is not None: + pattern = self.index_topk_pattern + if isinstance(pattern, str): + self.indexer_types = [ + {"F": "full", "S": "shared"}[c] for c in pattern + ] + else: + self.indexer_types = list(pattern) + else: + freq = max(self.index_topk_freq, 1) + offset = self.index_skip_topk_offset + self.indexer_types = [ + "full" if (max(i - offset + 1, 0) % freq) == 0 else "shared" + for i in range(self.num_hidden_layers) + ] + + +class GlmMoeDsaAttention(DeepseekV32Attention): + def __init__(self, config: ModelArgs, layer_idx: int): + super().__init__(config) + self.skip_topk = config.indexer_types[layer_idx] == "shared" + if self.skip_topk: + self.indexer = None + + def __call__( + self, + x: mx.array, + mask: Optional[mx.array] = None, + cache: Optional[Any] = None, + prev_topk_indices: Optional[mx.array] = None, + ): + B, L, D = x.shape + + qr = self.q_a_layernorm(self.q_a_proj(x)) + q = self.q_b_proj(qr) + + q = q.reshape(B, L, self.num_heads, self.q_head_dim).transpose(0, 2, 1, 3) + q_nope, q_pe = mx.split(q, [self.qk_nope_head_dim], axis=-1) + compressed_kv = self.kv_a_proj_with_mqa(x) + compressed_kv, k_pe = mx.split(compressed_kv, [self.kv_lora_rank], axis=-1) + k_pe = k_pe.reshape(B, L, 1, self.qk_rope_head_dim).transpose(0, 2, 1, 3) + kv_latent = self.kv_a_layernorm(compressed_kv) + + offset = cache[0].offset if cache is not None else 0 + q_pe = self.rope(q_pe, offset) + k_pe = self.rope(k_pe, offset) + + kv_latent = mx.expand_dims(kv_latent, axis=1) + + if cache is not None: + kv_latent, k_pe = cache[0].update_and_fetch(kv_latent, k_pe) + else: + cache = [None] * 2 + + if self.indexer is not None: + topk_indices = self.indexer(x, qr, mask, cache=cache[1]) + else: + topk_indices = prev_topk_indices + + if topk_indices is not None: + if L == 1: + idx = topk_indices[:, :, 0, :, None] + kv_latent = mx.take_along_axis( + kv_latent, + mx.broadcast_to(idx, idx.shape[:-1] + (kv_latent.shape[-1],)), + axis=2, + ) + k_pe = mx.take_along_axis( + k_pe, + mx.broadcast_to(idx, idx.shape[:-1] + (k_pe.shape[-1],)), + axis=2, + ) + if mask is not None: + mask = mx.take_along_axis(mask, topk_indices, axis=-1) + else: + shape = list(topk_indices.shape) + shape[-1] = kv_latent.shape[2] + sparse_mask = mx.zeros(shape, dtype=mx.bool_) + sparse_mask = mx.put_along_axis( + sparse_mask, topk_indices, mx.array(True), axis=-1 + ) + if mask is not None: + sparse_mask = sparse_mask & mask + mask = sparse_mask + + # Ensure the indexer cache is evaluated even if the topk_indices are unused + # to keep the graph from getting too large + if self.indexer is not None and cache is not None and cache[0] is not None: + cache[0].keys = mx.depends(cache[0].keys, (cache[1].keys, cache[1].values)) + + pe_scores = (q_pe * self.scale) @ k_pe.swapaxes(-1, -2) + if mask is not None: + pe_scores = mx.where( + mask, + pe_scores, + mx.array(mx.finfo(pe_scores.dtype).min, pe_scores.dtype), + ) + + if L == 1: + q_nope = self.embed_q(q_nope) + k = v = kv_latent + else: + k = self.embed_q(kv_latent, transpose=False) + v = self.unembed_out(kv_latent) + + output = scaled_dot_product_attention( + q_nope, k, v, cache=cache, scale=self.scale, mask=pe_scores + ) + if L == 1: + output = self.unembed_out(output) + + output = output.transpose(0, 2, 1, 3).reshape(B, L, -1) + return self.o_proj(output), topk_indices + + +class GlmMoeDsaDecoderLayer(DeepseekV32DecoderLayer): + def __init__(self, config: ModelArgs, layer_idx: int): + super().__init__(config, layer_idx) + self.self_attn = GlmMoeDsaAttention(config, layer_idx) + + def __call__( + self, + x: mx.array, + mask: Optional[mx.array] = None, + cache: Optional[Any] = None, + prev_topk_indices: Optional[mx.array] = None, + ): + r, topk_indices = self.self_attn( + self.input_layernorm(x), mask, cache, prev_topk_indices + ) + h = x + r + r = self.mlp(self.post_attention_layernorm(h)) + return h + r, topk_indices + + +class GlmMoeDsaModel(DeepseekV32Model): + def __init__(self, config: ModelArgs): + super().__init__(config) + self.layers = [ + GlmMoeDsaDecoderLayer(config, idx) + for idx in range(config.num_hidden_layers) + ] + + def __call__( + self, + x: mx.array, + cache: Optional[Any] = None, + ) -> mx.array: + h = self.embed_tokens(x) + + pipeline_rank = self.pipeline_rank + pipeline_size = self.pipeline_size + + if cache is None: + cache = [None] * self.num_layers + mask = create_attention_mask( + h, cache[0][0] if cache[0] else None, return_array=True + ) + + # Receive from the previous process in the pipeline + if pipeline_rank < pipeline_size - 1: + h = mx.distributed.recv_like(h, (pipeline_rank + 1)) + + prev_topk_indices = None + for i in range(self.num_layers): + h, prev_topk_indices = self.layers[self.start_idx + i]( + h, mask, cache[i], prev_topk_indices + ) + + # Send to the next process in the pipeline + if pipeline_rank != 0: + h = mx.distributed.send(h, (pipeline_rank - 1) % pipeline_size) + if cache[-1] is not None: + cache[-1][0].keys = mx.depends(cache[-1][0].keys, h) + + # Broadcast h while keeping it in the graph + if pipeline_size > 1: + h = mx.distributed.all_gather(h)[: h.shape[0]] + + return self.norm(h) + class Model(DSV32Model): def __init__(self, config: ModelArgs): super().__init__(config) + self.model = GlmMoeDsaModel(config) + + def make_cache(self): + # Shared layers run no indexer, so they get no indexer KVCache. + caches = [] + for layer in self.layers: + if getattr(layer.self_attn, "skip_topk", False): + caches.append(CacheList(KVCache())) + else: + caches.append(CacheList(KVCache(), KVCache())) + return caches diff --git a/tests/test_models.py b/tests/test_models.py index 6e1fcd96e..9b27c3e8d 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1422,6 +1422,66 @@ def test_deepseek_v3(self): model, args.model_type, args.vocab_size, args.num_hidden_layers ) + def test_glm_moe_dsa(self): + from mlx_lm.models import glm_moe_dsa + + args = glm_moe_dsa.ModelArgs( + model_type="glm_moe_dsa", + vocab_size=1024, + hidden_size=128, + index_head_dim=16, + index_n_heads=4, + index_topk=4, + intermediate_size=256, + moe_intermediate_size=256, + num_hidden_layers=6, + num_attention_heads=4, + num_key_value_heads=4, + n_shared_experts=1, + n_routed_experts=4, + routed_scaling_factor=2.5, + kv_lora_rank=16, + q_lora_rank=24, + qk_rope_head_dim=16, + v_head_dim=32, + qk_nope_head_dim=16, + topk_method="noaux_tc", + scoring_func="sigmoid", + norm_topk_prob=True, + n_group=2, + topk_group=1, + num_experts_per_tok=2, + moe_layer_freq=1, + first_k_dense_replace=1, + max_position_embeddings=1024, + rms_norm_eps=1e-5, + rope_parameters={"rope_theta": 10000.0}, + attention_bias=False, + index_topk_pattern="FSFSFS", + ) + self.assertEqual( + args.indexer_types, + ["full", "shared", "full", "shared", "full", "shared"], + ) + model = glm_moe_dsa.Model(args) + + has_indexer = [l.self_attn.indexer is not None for l in model.model.layers] + self.assertEqual(has_indexer, [True, False, True, False, True, False]) + + self.model_test_runner( + model, args.model_type, args.vocab_size, args.num_hidden_layers + ) + + prompt = mx.array([[1, 2, 3, 4, 5, 6, 7, 8]]) + cache = make_prompt_cache(model) + logits = model(prompt, cache=cache) + self.assertEqual(logits.shape, (1, 8, args.vocab_size)) + nxt = mx.argmax(logits[0, -1:, :], keepdims=True) + logits = model(nxt, cache=cache) + self.assertEqual(logits.shape, (1, 1, args.vocab_size)) + self.assertTrue(mx.all(mx.isfinite(logits)).item()) + mx.eval([c.state for c in cache]) + def test_gemma2(self): from mlx_lm.models import gemma2