Support per-layer DSA indexer schedule (GLM-5.2)

mlx_lm/models/deepseek_v32.py

@@ -114,7 +114,7 @@ def __call__(
 
 
 class DeepseekV32Attention(nn.Module):
-    def __init__(self, config: ModelArgs):
+    def __init__(self, config: ModelArgs, layer_idx: int = None):
         super().__init__()
         self.config = config
         self.hidden_size = config.hidden_size
@@ -165,7 +165,15 @@ def __init__(self, config: ModelArgs):
                     s = 0.1 * mscale_all_dim * math.log(scaling_factor) + 1.0
                     self.scale = self.scale * s * s
 
-        self.indexer = Indexer(config)
+        # Some DSA models (e.g. GLM-5.2) enable the lightning indexer on only a
+        # subset of layers via a per-layer `indexer_types` list ("full" = own
+        # indexer, otherwise none). When the field is absent (DeepSeek-V3.2) the
+        # indexer is built on every layer, as before.
+        indexer_types = getattr(config, "indexer_types", None)
+        if indexer_types is None or layer_idx is None or indexer_types[layer_idx] == "full":
+            self.indexer = Indexer(config)
+        else:
+            self.indexer = None
         self.rope = initialize_rope(
             dims=self.qk_rope_head_dim,
             base=self.rope_theta,
@@ -203,36 +211,37 @@ def __call__(
         else:
             cache = [None] * 2
 
-        topk_indices = self.indexer(x, qr, mask, cache=cache[1])
-        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 cache is not None and cache[0] is not None:
-            cache[0].keys = mx.depends(cache[0].keys, (cache[1].keys, cache[1].values))
+        if self.indexer is not None:
+            topk_indices = self.indexer(x, qr, mask, cache=cache[1])
+            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 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:
@@ -379,7 +388,7 @@ def __call__(self, x):
 class DeepseekV32DecoderLayer(nn.Module):
     def __init__(self, config: ModelArgs, layer_idx: int):
         super().__init__()
-        self.self_attn = DeepseekV32Attention(config)
+        self.self_attn = DeepseekV32Attention(config, layer_idx)
         self.mlp = (
             DeepseekV32MoE(config)
             if (
@@ -651,4 +660,13 @@ def predicate(k):
         return predicate
 
     def make_cache(self):
-        return [CacheList(KVCache(), KVCache()) for _ in self.layers]
+        # Layers without an indexer (DSA disabled for that layer) need only the
+        # main attention cache; a second, never-written KVCache would crash
+        # CacheList.state on its empty `keys`.
+        caches = []
+        for layer in self.layers:
+            if layer is not None and getattr(layer.self_attn, "indexer", None) is None:
+                caches.append(CacheList(KVCache()))
+            else:
+                caches.append(CacheList(KVCache(), KVCache()))
+        return caches

mlx_lm/models/glm_moe_dsa.py

@@ -42,6 +42,9 @@ class ModelArgs(BaseModelArgs):
     attention_bias: bool
     rope_scaling: Dict = None
     rope_theta: Optional[float] = None
+    # Per-layer DSA indexer schedule ("full" -> own indexer, else none). GLM-5.2
+    # uses this; absent -> indexer on every layer (DeepSeek-V3.2 behaviour).
+    indexer_types: Optional[Any] = None
 
     def __post_init__(self):
         self.rope_scaling = self.rope_parameters

tests/test_models.py

@@ -2888,6 +2888,42 @@ def test_all_models(self):
                     "type": "yarn",
                 },
             },
+            {
+                # GLM-5.2 style: DSA indexer on only a subset of layers via
+                # `indexer_types` ("full" -> own indexer, else none).
+                "model_type": "glm_moe_dsa",
+                "vocab_size": 1024,
+                "hidden_size": 128,
+                "index_head_dim": 32,
+                "index_n_heads": 2,
+                "index_topk": 2048,
+                "intermediate_size": 256,
+                "moe_intermediate_size": 256,
+                "num_hidden_layers": 4,
+                "num_attention_heads": 4,
+                "num_key_value_heads": 2,
+                "n_shared_experts": 1,
+                "n_routed_experts": 4,
+                "routed_scaling_factor": 1.0,
+                "kv_lora_rank": 4,
+                "q_lora_rank": 4,
+                "qk_rope_head_dim": 32,
+                "v_head_dim": 16,
+                "qk_nope_head_dim": 32,
+                "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": 4096,
+                "rms_norm_eps": 1e-6,
+                "rope_parameters": {"rope_theta": 10000, "rope_type": "default"},
+                "attention_bias": False,
+                "indexer_types": ["full", "shared", "full", "shared"],
+            },
             {
                 "model_type": "mimo_v2_flash",
                 "num_experts_per_tok": 2,