diff --git a/mlx_lm/models/deepseek_v32.py b/mlx_lm/models/deepseek_v32.py index 7c97682e7..4e9c13cd7 100644 --- a/mlx_lm/models/deepseek_v32.py +++ b/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 diff --git a/mlx_lm/models/glm_moe_dsa.py b/mlx_lm/models/glm_moe_dsa.py index 14e96e365..738100bba 100644 --- a/mlx_lm/models/glm_moe_dsa.py +++ b/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 diff --git a/tests/test_models.py b/tests/test_models.py index 6e1fcd96e..5f4db2f75 100644 --- a/tests/test_models.py +++ b/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,