[transformers-to-mlx skill] glm_moe_dsa: DSA cross-layer indexer sharing (GLM-5.2)

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
 
 
@@ -43,12 +51,209 @@ class ModelArgs(BaseModelArgs):
     rope_scaling: Dict = None
     rope_theta: Optional[float] = None
     indexer_rope_interleave: bool = True
+    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

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