From 2bfc36155b2edcb38a261113fc5077b48bb9dd8f Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Wed, 12 Nov 2025 12:45:37 +0100 Subject: [PATCH 1/8] fix input_embeddings prefill bug in generate_step --- mlx_lm/generate.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/mlx_lm/generate.py b/mlx_lm/generate.py index 6ca8a73fb..e9d48e74e 100644 --- a/mlx_lm/generate.py +++ b/mlx_lm/generate.py @@ -418,7 +418,8 @@ def _step(input_tokens: mx.array, input_embeddings: Optional[mx.array] = None): prompt_processed_tokens = 0 prompt_progress_callback(prompt_processed_tokens, total_prompt_tokens) while total_prompt_tokens - prompt_processed_tokens > 1: - n_to_process = min(prefill_step_size, prompt.size - 1) + remaining = (total_prompt_tokens - prompt_processed_tokens) - 1 + n_to_process = min(prefill_step_size, remaining) _model_call( input_tokens=prompt[:n_to_process][None], input_embeddings=( From 5cf134cae379fe3b5771356d18d5d8b44deb9481 Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Wed, 12 Nov 2025 12:56:42 +0100 Subject: [PATCH 2/8] format --- mlx_lm/generate.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlx_lm/generate.py b/mlx_lm/generate.py index e9d48e74e..1dda46c2f 100644 --- a/mlx_lm/generate.py +++ b/mlx_lm/generate.py @@ -418,7 +418,7 @@ def _step(input_tokens: mx.array, input_embeddings: Optional[mx.array] = None): prompt_processed_tokens = 0 prompt_progress_callback(prompt_processed_tokens, total_prompt_tokens) while total_prompt_tokens - prompt_processed_tokens > 1: - remaining = (total_prompt_tokens - prompt_processed_tokens) - 1 + remaining = (total_prompt_tokens - prompt_processed_tokens) - 1 n_to_process = min(prefill_step_size, remaining) _model_call( input_tokens=prompt[:n_to_process][None], From e52be25843a691d44089464e73128be06ab650d4 Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Mon, 27 Apr 2026 00:02:56 +0200 Subject: [PATCH 3/8] add laguna --- mlx_lm/models/laguna.py | 461 ++++++++++++++++++++++++++++++++++++++++ tests/test_models.py | 103 +++++++++ 2 files changed, 564 insertions(+) create mode 100644 mlx_lm/models/laguna.py diff --git a/mlx_lm/models/laguna.py b/mlx_lm/models/laguna.py new file mode 100644 index 000000000..2d8f3cf69 --- /dev/null +++ b/mlx_lm/models/laguna.py @@ -0,0 +1,461 @@ +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional, Union + +import mlx.core as mx +import mlx.nn as nn + +from .activations import swiglu +from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention +from .rope_utils import initialize_rope +from .switch_layers import SwitchGLU + + +@dataclass +class ModelArgs(BaseModelArgs): + model_type: str + vocab_size: int + hidden_size: int + intermediate_size: int + num_hidden_layers: int + num_attention_heads: int + num_key_value_heads: int + head_dim: int + max_position_embeddings: int + rms_norm_eps: float = 1e-6 + qkv_bias: bool = False + attention_bias: bool = False + gating: Union[bool, str] = True + tie_word_embeddings: bool = False + rope_theta: float = 500000.0 + rope_parameters: Optional[Dict[str, Union[float, str]]] = None + rope_scaling: Optional[Dict[str, Union[float, str]]] = None + partial_rotary_factor: Optional[float] = None + rope_style: str = "rotate-half" + sliding_window: Optional[int] = None + layer_types: Optional[List[str]] = None + num_attention_heads_per_layer: Optional[List[int]] = None + swa_rope_parameters: Optional[Dict[str, Union[float, str]]] = None + swa_attention_sink_enabled: bool = False + num_experts: int = 0 + num_experts_per_tok: int = 0 + moe_intermediate_size: int = 0 + shared_expert_intermediate_size: int = 0 + norm_topk_prob: bool = True + decoder_sparse_step: int = 1 + mlp_only_layers: List[int] = field(default_factory=lambda: [0]) + moe_routed_scaling_factor: float = 1.0 + moe_apply_router_weight_on_input: bool = False + moe_router_logit_softcapping: float = 0.0 + moe_router_use_sigmoid: bool = True + + def __post_init__(self): + if self.layer_types is None: + self.layer_types = ["full_attention"] * self.num_hidden_layers + if len(self.layer_types) != self.num_hidden_layers: + raise ValueError("layer_types must match num_hidden_layers.") + + if self.num_attention_heads_per_layer is None: + self.num_attention_heads_per_layer = [ + self.num_attention_heads + ] * self.num_hidden_layers + if len(self.num_attention_heads_per_layer) != self.num_hidden_layers: + raise ValueError( + "num_attention_heads_per_layer must match num_hidden_layers." + ) + if any( + h % self.num_key_value_heads for h in self.num_attention_heads_per_layer + ): + raise ValueError( + "Every query-head count must be divisible by num_key_value_heads." + ) + + if self.rope_parameters is None: + self.rope_parameters = ( + dict(self.rope_scaling) + if self.rope_scaling is not None + else {"rope_type": "default", "rope_theta": self.rope_theta} + ) + else: + self.rope_parameters = dict(self.rope_parameters) + + if self.swa_rope_parameters is not None: + self.swa_rope_parameters = dict(self.swa_rope_parameters) + + if self.partial_rotary_factor is not None: + self.rope_parameters.setdefault( + "partial_rotary_factor", self.partial_rotary_factor + ) + if self.swa_rope_parameters is not None: + self.swa_rope_parameters.setdefault( + "partial_rotary_factor", self.partial_rotary_factor + ) + + +def _rope_base(args: ModelArgs, rope_config: Dict[str, Union[float, str]]) -> float: + return float(rope_config.get("rope_theta", args.rope_theta)) + + +def _rope_dims(args: ModelArgs, rope_config: Dict[str, Union[float, str]]) -> int: + partial = float(rope_config.get("partial_rotary_factor", 1.0)) + return int(args.head_dim * partial) + + +class MLP(nn.Module): + def __init__(self, dim: int, hidden_dim: int): + super().__init__() + self.gate_proj = nn.Linear(dim, hidden_dim, bias=False) + self.down_proj = nn.Linear(hidden_dim, dim, bias=False) + self.up_proj = nn.Linear(dim, hidden_dim, bias=False) + + def __call__(self, x) -> mx.array: + return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x))) + + +class LagunaTopKRouter(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.top_k = args.num_experts_per_tok + self.norm_topk_prob = args.norm_topk_prob + self.use_sigmoid = args.moe_router_use_sigmoid + self.router_logit_softcapping = args.moe_router_logit_softcapping + self.proj = nn.Linear(args.hidden_size, args.num_experts, bias=False) + self.e_score_correction_bias = mx.zeros((args.num_experts,)) + + def __call__(self, x: mx.array) -> tuple[mx.array, mx.array]: + dtype = x.dtype + logits = self.proj(x).astype(mx.float32) + if self.router_logit_softcapping > 0.0: + c = self.router_logit_softcapping + logits = mx.tanh(logits / c) * c + + scores = mx.sigmoid(logits) if self.use_sigmoid else mx.softmax(logits, axis=-1) + corrected_scores = scores + self.e_score_correction_bias.astype(scores.dtype) + + k = self.top_k + inds = mx.stop_gradient( + mx.argpartition(-corrected_scores, kth=k - 1, axis=-1)[..., :k] + ) + weights = mx.take_along_axis(scores, inds, axis=-1) + if self.norm_topk_prob: + weights = weights / mx.sum(weights, axis=-1, keepdims=True) + return inds, weights.astype(dtype) + + +class LagunaSparseMoeBlock(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + if args.moe_apply_router_weight_on_input: + raise NotImplementedError( + "moe_apply_router_weight_on_input=True is not supported." + ) + self.routed_scaling_factor = args.moe_routed_scaling_factor + self.gate = LagunaTopKRouter(args) + self.switch_mlp = SwitchGLU( + args.hidden_size, args.moe_intermediate_size, args.num_experts + ) + self.shared_expert = MLP(args.hidden_size, args.shared_expert_intermediate_size) + + def __call__(self, x: mx.array) -> mx.array: + inds, scores = self.gate(x) + y = self.switch_mlp(x, inds) + y = mx.sum(y * scores[..., None], axis=-2) + if self.routed_scaling_factor != 1.0: + y = y * self.routed_scaling_factor + return y + self.shared_expert(x) + + +class Attention(nn.Module): + def __init__(self, args: ModelArgs, layer_idx: int): + super().__init__() + + self.n_heads = args.num_attention_heads_per_layer[layer_idx] + self.n_kv_heads = args.num_key_value_heads + self.head_dim = args.head_dim + self.scale = self.head_dim**-0.5 + self.gate_per_head = args.gating == "per-head" + self.gating = bool(args.gating) + self.is_sliding = args.layer_types[layer_idx] == "sliding_attention" + self.sliding_window = args.sliding_window if self.is_sliding else None + + dim = args.hidden_size + self.q_proj = nn.Linear( + dim, self.n_heads * self.head_dim, bias=args.qkv_bias + ) + self.k_proj = nn.Linear( + dim, self.n_kv_heads * self.head_dim, bias=args.qkv_bias + ) + self.v_proj = nn.Linear( + dim, self.n_kv_heads * self.head_dim, bias=args.qkv_bias + ) + self.o_proj = nn.Linear( + self.n_heads * self.head_dim, dim, bias=args.attention_bias + ) + + if self.gating: + gate_dim = ( + self.n_heads + if self.gate_per_head + else self.n_heads * self.head_dim + ) + self.g_proj = nn.Linear(dim, gate_dim, bias=False) + + if self.is_sliding and args.swa_attention_sink_enabled: + self.sink = mx.zeros((self.n_heads,)) + else: + self.sink = None + + self.q_norm = nn.RMSNorm(self.head_dim, eps=args.rms_norm_eps) + self.k_norm = nn.RMSNorm(self.head_dim, eps=args.rms_norm_eps) + + rope_config = ( + args.swa_rope_parameters + if self.is_sliding and args.swa_rope_parameters is not None + else args.rope_parameters + ) + self.rope = initialize_rope( + _rope_dims(args, rope_config), + base=_rope_base(args, rope_config), + traditional=False, + scaling_config=rope_config, + max_position_embeddings=args.max_position_embeddings, + ) + + def __call__( + self, + x: mx.array, + mask: Optional[mx.array] = None, + cache: Optional[Any] = None, + ) -> mx.array: + B, L, _ = x.shape + + queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x) + queries = self.q_norm( + queries.reshape(B, L, self.n_heads, self.head_dim) + ).transpose(0, 2, 1, 3) + keys = self.k_norm( + keys.reshape(B, L, self.n_kv_heads, self.head_dim) + ).transpose(0, 2, 1, 3) + values = values.reshape(B, L, self.n_kv_heads, self.head_dim).transpose( + 0, 2, 1, 3 + ) + + if cache is not None: + queries = self.rope(queries, offset=cache.offset) + keys = self.rope(keys, offset=cache.offset) + keys, values = cache.update_and_fetch(keys, values) + else: + queries = self.rope(queries) + keys = self.rope(keys) + + output = scaled_dot_product_attention( + queries, + keys, + values, + cache=cache, + scale=self.scale, + mask=mask, + sinks=self.sink, + ) + output = output.transpose(0, 2, 1, 3).reshape(B, L, -1) + + if self.gating: + gate = nn.softplus(self.g_proj(x).astype(mx.float32)).astype(output.dtype) + if self.gate_per_head: + shape = output.shape + output = ( + output.reshape(B, L, self.n_heads, self.head_dim) + * gate[..., None] + ).reshape(shape) + else: + output = output * gate + + return self.o_proj(output) + + +class DecoderLayer(nn.Module): + def __init__(self, args: ModelArgs, layer_idx: int): + super().__init__() + self.self_attn = Attention(args, layer_idx) + if (layer_idx not in args.mlp_only_layers) and ( + args.num_experts > 0 and (layer_idx + 1) % args.decoder_sparse_step == 0 + ): + self.mlp = LagunaSparseMoeBlock(args) + else: + self.mlp = MLP(args.hidden_size, args.intermediate_size) + + self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps) + self.post_attention_layernorm = nn.RMSNorm( + args.hidden_size, eps=args.rms_norm_eps + ) + self.attention_type = args.layer_types[layer_idx] + + def __call__( + self, + x: mx.array, + mask: Optional[mx.array] = None, + cache: Optional[Any] = None, + ) -> mx.array: + r = self.self_attn(self.input_layernorm(x), mask, cache) + h = x + r + r = self.mlp(self.post_attention_layernorm(h)) + return h + r + + +class LagunaModel(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.args = args + self.vocab_size = args.vocab_size + self.num_hidden_layers = args.num_hidden_layers + self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size) + self.layers = [ + DecoderLayer(args, layer_idx) for layer_idx in range(args.num_hidden_layers) + ] + self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps) + self.fa_idx = args.layer_types.index("full_attention") + self.swa_idx = ( + args.layer_types.index("sliding_attention") + if "sliding_attention" in args.layer_types + else None + ) + + def __call__( + self, + inputs: mx.array, + cache=None, + input_embeddings: Optional[mx.array] = None, + ) -> mx.array: + if input_embeddings is not None: + h = input_embeddings + else: + h = self.embed_tokens(inputs) + + if cache is None: + cache = [None] * len(self.layers) + + full_mask = create_attention_mask(h, cache[self.fa_idx]) + if self.swa_idx is not None: + sliding_mask = create_attention_mask( + h, cache[self.swa_idx], window_size=self.args.sliding_window + ) + + for layer, c in zip(self.layers, cache): + mask = ( + sliding_mask + if layer.attention_type == "sliding_attention" + else full_mask + ) + h = layer(h, mask, c) + return self.norm(h) + + +class Model(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.args = args + self.model_type = args.model_type + self.model = LagunaModel(args) + if not args.tie_word_embeddings: + self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False) + + def __call__( + self, + inputs: mx.array, + cache=None, + input_embeddings: Optional[mx.array] = None, + ) -> mx.array: + out = self.model(inputs, cache, input_embeddings) + if self.args.tie_word_embeddings: + return self.model.embed_tokens.as_linear(out) + return self.lm_head(out) + + def sanitize(self, weights): + if self.args.tie_word_embeddings: + weights.pop("lm_head.weight", None) + + weights = self._unpack_compressed_tensors(weights) + weights = self._remap_router_weights(weights) + weights = self._stack_experts(weights) + return { + k: v + for k, v in weights.items() + if "rotary_emb.inv_freq" not in k + and not k.endswith(".self_attn.k_scale") + and not k.endswith(".self_attn.v_scale") + } + + def _unpack_compressed_tensors(self, weights): + if not any(k.endswith(".weight_shape") for k in weights): + return weights + + new_weights = {} + for k, v in weights.items(): + if k.endswith(".weight_shape"): + base = k[: -len("weight_shape")] + if ( + f"{base}weight_packed" in weights + and f"{base}weight_scale" in weights + ): + scales = weights[f"{base}weight_scale"] + new_weights[f"{base}weight"] = weights[ + f"{base}weight_packed" + ].view(mx.uint32) + new_weights[f"{base}scales"] = scales + new_weights[f"{base}biases"] = (-8 * scales).astype(scales.dtype) + elif k.endswith(".weight_packed") or k.endswith(".weight_scale"): + base = k.rsplit(".", 1)[0] + "." + if f"{base}weight_shape" in weights: + continue + new_weights[k] = v + else: + new_weights[k] = v + return new_weights + + def _remap_router_weights(self, weights): + for layer_idx in range(self.args.num_hidden_layers): + prefix = f"model.layers.{layer_idx}.mlp" + gate_weight = f"{prefix}.gate.weight" + if gate_weight in weights: + weights[f"{prefix}.gate.proj.weight"] = weights.pop(gate_weight) + + legacy_bias = f"{prefix}.experts.e_score_correction_bias" + if legacy_bias in weights: + weights[f"{prefix}.gate.e_score_correction_bias"] = weights.pop( + legacy_bias + ) + return weights + + def _stack_experts(self, weights): + for layer_idx in range(self.args.num_hidden_layers): + prefix = f"model.layers.{layer_idx}.mlp" + for proj in ["gate_proj", "up_proj", "down_proj"]: + for suffix in ["weight", "scales", "biases"]: + first_key = f"{prefix}.experts.0.{proj}.{suffix}" + if first_key not in weights: + continue + weights[f"{prefix}.switch_mlp.{proj}.{suffix}"] = mx.stack( + [ + weights.pop(f"{prefix}.experts.{e}.{proj}.{suffix}") + for e in range(self.args.num_experts) + ] + ) + return weights + + @property + def quant_predicate(self): + def predicate(path, _): + if path.endswith("mlp.gate.proj"): + return {"group_size": 64, "bits": 8} + return True + + return predicate + + @property + def cast_predicate(self): + def predicate(k): + return "e_score_correction_bias" not in k + return predicate + + @property + def layers(self): + return self.model.layers diff --git a/tests/test_models.py b/tests/test_models.py index 6e1fcd96e..70cc82ee7 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -842,6 +842,68 @@ def test_gemma4_moe_router_quantizes_to_8bit(self): ) self.assertEqual(config["quantization"]["bits"], 4) + def test_laguna_sanitize_hf_moe_weights(self): + from mlx_lm.models import laguna + + args = laguna.ModelArgs.from_dict( + { + "model_type": "laguna", + "vocab_size": 32, + "hidden_size": 8, + "intermediate_size": 16, + "num_hidden_layers": 2, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "head_dim": 4, + "max_position_embeddings": 32, + "num_experts": 2, + "num_experts_per_tok": 1, + "decoder_sparse_step": 1, + "mlp_only_layers": [0], + "moe_intermediate_size": 3, + "shared_expert_intermediate_size": 3, + "layer_types": ["full_attention", "sliding_attention"], + "sliding_window": 4, + "swa_rope_parameters": { + "rope_theta": 10000.0, + "rope_type": "linear", + "factor": 1.0, + "partial_rotary_factor": 1.0, + }, + } + ) + model = laguna.Model(args) + + weights = { + "model.layers.1.mlp.gate.weight": mx.ones((2, 8)), + "model.layers.1.mlp.experts.e_score_correction_bias": mx.arange( + 2, dtype=mx.float32 + ), + } + for e in range(2): + weights[f"model.layers.1.mlp.experts.{e}.gate_proj.weight"] = ( + mx.full((3, 8), e + 1, dtype=mx.float32) + ) + weights[f"model.layers.1.mlp.experts.{e}.up_proj.weight"] = mx.full( + (3, 8), e + 3, dtype=mx.float32 + ) + weights[f"model.layers.1.mlp.experts.{e}.down_proj.weight"] = mx.full( + (8, 3), e + 5, dtype=mx.float32 + ) + + converted = model.sanitize(weights) + + self.assertIn("model.layers.1.mlp.gate.proj.weight", converted) + self.assertIn("model.layers.1.mlp.gate.e_score_correction_bias", converted) + self.assertIn("model.layers.1.mlp.switch_mlp.gate_proj.weight", converted) + self.assertIn("model.layers.1.mlp.switch_mlp.up_proj.weight", converted) + self.assertIn("model.layers.1.mlp.switch_mlp.down_proj.weight", converted) + self.assertEqual( + converted["model.layers.1.mlp.switch_mlp.gate_proj.weight"].shape, + (2, 3, 8), + ) + self.assertFalse(any(".experts." in k for k in converted)) + def test_qwen2_moe(self): from mlx_lm.models import qwen2_moe @@ -2585,6 +2647,47 @@ def test_all_models(self): "max_position_embeddings": 1000, "norm_topk_prob": True, }, + { + "model_type": "laguna", + "vocab_size": 1000, + "hidden_size": 128, + "intermediate_size": 256, + "num_hidden_layers": 4, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "head_dim": 16, + "max_position_embeddings": 1000, + "rms_norm_eps": 1e-5, + "num_experts": 4, + "num_experts_per_tok": 2, + "decoder_sparse_step": 1, + "mlp_only_layers": [0], + "moe_intermediate_size": 64, + "shared_expert_intermediate_size": 64, + "norm_topk_prob": True, + "moe_routed_scaling_factor": 2.5, + "gating": "per-head", + "sliding_window": 4, + "layer_types": [ + "full_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + ], + "num_attention_heads_per_layer": [4, 8, 8, 4], + "rope_parameters": { + "rope_theta": 10000.0, + "rope_type": "default", + "partial_rotary_factor": 0.5, + }, + "swa_rope_parameters": { + "rope_theta": 10000.0, + "rope_type": "linear", + "factor": 1.0, + "partial_rotary_factor": 1.0, + }, + "tie_word_embeddings": False, + }, { "model_type": "lille-130m", "block_size": 128, From b50cc73014c7e98113b3b34878cda46016c35586 Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Mon, 27 Apr 2026 00:20:56 +0200 Subject: [PATCH 4/8] refactor compressed tensors to load INT4 quants --- mlx_lm/utils.py | 27 ++++++++++++++++++++++++++- tests/test_utils.py | 17 +++++++++++++++++ 2 files changed, 43 insertions(+), 1 deletion(-) diff --git a/mlx_lm/utils.py b/mlx_lm/utils.py index ef3d266b9..93b5758ef 100644 --- a/mlx_lm/utils.py +++ b/mlx_lm/utils.py @@ -172,6 +172,31 @@ def _transform_awq_weights( return new_weights, mlx_quantization +def _compressed_tensors_config(quantization_config: Dict[str, Any]) -> Dict[str, Any]: + config_groups = quantization_config.get("config_groups") + if not config_groups: + # Default quantization settings + return {"group_size": 32, "bits": 4, "mode": "affine"} + + # Use first config group for parameters + first_group = next(iter(config_groups.values())) + weights_config = first_group.get("weights", {}) + + format_type = quantization_config.get("format") + weights_type = weights_config.get("type") + + if format_type == "pack-quantized" and weights_type == "int": + return { + "group_size": weights_config.get("group_size", 32), + "bits": weights_config.get("num_bits", 4), + "mode": "affine", + } + + raise ValueError( + f"Unsupported compressed-tensors quantization format: {format_type}" + ) + + def _get_classes(config: dict): """ Retrieve the model and model args classes based on the configuration. @@ -378,7 +403,7 @@ def class_predicate(p, m): config["quantization_config"] = quantization _quantize(quantization) elif quant_method == "compressed-tensors": - quantization = {"group_size": 32, "bits": 4, "mode": "affine"} + quantization = _compressed_tensors_config(quantization_config) config["quantization"] = quantization config["quantization_config"] = quantization _quantize(quantization) diff --git a/tests/test_utils.py b/tests/test_utils.py index 88b68fe33..e7f19abe3 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -79,6 +79,23 @@ def test_quantize(self): self.assertEqual(config["quantization"]["group_size"], 64) self.assertEqual(config["quantization"]["bits"], 4) + def test_compressed_tensors_config_uses_weight_group_size(self): + quantization = utils._compressed_tensors_config( + { + "format": "pack-quantized", + "config_groups": { + "group_0": { + "weights": { + "type": "int", + "num_bits": 4, + "group_size": 128, + } + } + }, + } + ) + self.assertEqual(quantization, {"group_size": 128, "bits": 4, "mode": "affine"}) + def test_convert(self): mlx_path = os.path.join(self.test_dir, "mlx_model") From 4264da550e2fa87b09430cd8f586c2fb28ef3160 Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Mon, 27 Apr 2026 01:50:52 +0200 Subject: [PATCH 5/8] add tool parsing --- mlx_lm/tokenizer_utils.py | 6 ++++++ mlx_lm/tool_parsers/laguna.py | 35 ++++++++++++++++++++++++++++++++++ tests/test_tool_parsing.py | 36 +++++++++++++++++++++++++++++++++++ 3 files changed, 77 insertions(+) create mode 100644 mlx_lm/tool_parsers/laguna.py diff --git a/mlx_lm/tokenizer_utils.py b/mlx_lm/tokenizer_utils.py index c7e50fbe7..f475e9f9b 100644 --- a/mlx_lm/tokenizer_utils.py +++ b/mlx_lm/tokenizer_utils.py @@ -557,6 +557,12 @@ def _infer_tool_parser(chat_template): return "function_gemma" elif "" in chat_template: return "longcat" + elif ( + "function-name" in chat_template + and "" in chat_template + and "" in chat_template + ): + return "laguna" elif "" in chat_template: return "glm47" elif "<|tool_list_start|>" in chat_template: diff --git a/mlx_lm/tool_parsers/laguna.py b/mlx_lm/tool_parsers/laguna.py new file mode 100644 index 000000000..8ac295a75 --- /dev/null +++ b/mlx_lm/tool_parsers/laguna.py @@ -0,0 +1,35 @@ +# Copyright © 2026 Apple Inc. + +from typing import Any + +from . import glm47 + +tool_call_start = "" +tool_call_end = "" + + +def _strip_names(tool_call): + if isinstance(tool_call, list): + return [_strip_names(tc) for tc in tool_call] + if isinstance(tool_call, dict) and isinstance(tool_call.get("name"), str): + tool_call = dict(tool_call) + tool_call["name"] = tool_call["name"].strip() + return tool_call + + +def parse_tool_call(text: str, tools: list[Any] | None = None): + """Parse Laguna XML-like tool calls into a name and arguments dict.""" + match = glm47._func_name_regex.search(text) + if not match: + return _strip_names(glm47.parse_tool_call(text, tools)) + + func_name = match.group(1).strip() + string_args = glm47._get_string_arg_names(func_name, tools) + arguments = {} + for match in glm47._func_arg_regex.finditer(text): + arg_key = match.group(1).strip() + arg_val = match.group(2).strip() + if arg_key not in string_args: + arg_val = glm47._deserialize(arg_val) + arguments[arg_key] = arg_val + return dict(name=func_name, arguments=arguments) diff --git a/tests/test_tool_parsing.py b/tests/test_tool_parsing.py index 52892b7ff..f134edafe 100644 --- a/tests/test_tool_parsing.py +++ b/tests/test_tool_parsing.py @@ -7,6 +7,7 @@ glm47, json_tools, kimi_k2, + laguna, longcat, minimax_m2, mistral, @@ -329,6 +330,41 @@ def test_minimax_m2(self): tool_calls = minimax_m2.parse_tool_call(test_case, None) self.assertEqual(expected, tool_calls) + def test_laguna_strips_function_name_and_preserves_string_args(self): + tools = [ + { + "type": "function", + "function": { + "name": "get_weather", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string"}, + "days": {"type": "integer"}, + "postal_code": {"type": "string"}, + }, + }, + }, + } + ] + test_case = ( + "get_weather\n" + "location\n" + "Warsaw\n" + "days\n" + "3\n" + "postal_code\n" + "00123" + ) + + tool_call = laguna.parse_tool_call(test_case, tools) + + self.assertEqual(tool_call["name"], "get_weather") + self.assertEqual( + tool_call["arguments"], + {"location": "Warsaw", "days": 3, "postal_code": "00123"}, + ) + if __name__ == "__main__": unittest.main() From b88b0a6e3d716b0962705bdfa00e47f6b7c88448 Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Mon, 27 Apr 2026 01:51:41 +0200 Subject: [PATCH 6/8] format --- mlx_lm/models/laguna.py | 18 +++++++----------- tests/test_models.py | 4 ++-- 2 files changed, 9 insertions(+), 13 deletions(-) diff --git a/mlx_lm/models/laguna.py b/mlx_lm/models/laguna.py index 2d8f3cf69..79ca4325d 100644 --- a/mlx_lm/models/laguna.py +++ b/mlx_lm/models/laguna.py @@ -178,9 +178,7 @@ def __init__(self, args: ModelArgs, layer_idx: int): self.sliding_window = args.sliding_window if self.is_sliding else None dim = args.hidden_size - self.q_proj = nn.Linear( - dim, self.n_heads * self.head_dim, bias=args.qkv_bias - ) + self.q_proj = nn.Linear(dim, self.n_heads * self.head_dim, bias=args.qkv_bias) self.k_proj = nn.Linear( dim, self.n_kv_heads * self.head_dim, bias=args.qkv_bias ) @@ -193,9 +191,7 @@ def __init__(self, args: ModelArgs, layer_idx: int): if self.gating: gate_dim = ( - self.n_heads - if self.gate_per_head - else self.n_heads * self.head_dim + self.n_heads if self.gate_per_head else self.n_heads * self.head_dim ) self.g_proj = nn.Linear(dim, gate_dim, bias=False) @@ -263,8 +259,7 @@ def __call__( if self.gate_per_head: shape = output.shape output = ( - output.reshape(B, L, self.n_heads, self.head_dim) - * gate[..., None] + output.reshape(B, L, self.n_heads, self.head_dim) * gate[..., None] ).reshape(shape) else: output = output * gate @@ -397,9 +392,9 @@ def _unpack_compressed_tensors(self, weights): and f"{base}weight_scale" in weights ): scales = weights[f"{base}weight_scale"] - new_weights[f"{base}weight"] = weights[ - f"{base}weight_packed" - ].view(mx.uint32) + new_weights[f"{base}weight"] = weights[f"{base}weight_packed"].view( + mx.uint32 + ) new_weights[f"{base}scales"] = scales new_weights[f"{base}biases"] = (-8 * scales).astype(scales.dtype) elif k.endswith(".weight_packed") or k.endswith(".weight_scale"): @@ -454,6 +449,7 @@ def predicate(path, _): def cast_predicate(self): def predicate(k): return "e_score_correction_bias" not in k + return predicate @property diff --git a/tests/test_models.py b/tests/test_models.py index 70cc82ee7..a9418ca1d 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -881,8 +881,8 @@ def test_laguna_sanitize_hf_moe_weights(self): ), } for e in range(2): - weights[f"model.layers.1.mlp.experts.{e}.gate_proj.weight"] = ( - mx.full((3, 8), e + 1, dtype=mx.float32) + weights[f"model.layers.1.mlp.experts.{e}.gate_proj.weight"] = mx.full( + (3, 8), e + 1, dtype=mx.float32 ) weights[f"model.layers.1.mlp.experts.{e}.up_proj.weight"] = mx.full( (3, 8), e + 3, dtype=mx.float32 From aa8eb90c6cb84850224635a9e30c8578c191e72b Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Tue, 28 Apr 2026 16:52:49 +0200 Subject: [PATCH 7/8] refactor tool parser and tests --- mlx_lm/tool_parsers/laguna.py | 80 ++++++++++++++++++++++++------- tests/test_models.py | 90 +++++++++++++++++++---------------- tests/test_tool_parsing.py | 46 ++++++++++++++++++ 3 files changed, 159 insertions(+), 57 deletions(-) diff --git a/mlx_lm/tool_parsers/laguna.py b/mlx_lm/tool_parsers/laguna.py index 8ac295a75..444b37f41 100644 --- a/mlx_lm/tool_parsers/laguna.py +++ b/mlx_lm/tool_parsers/laguna.py @@ -1,35 +1,83 @@ # Copyright © 2026 Apple Inc. +""" +Tool parser for Poolside Laguna XML-like tool calls. + +Format: +function-name +argument-key +value-of-argument-key + +""" + +import ast +import json from typing import Any -from . import glm47 +import regex as re tool_call_start = "" tool_call_end = "" +_tool_call_regex = re.compile(r"(.*?)", re.DOTALL) +_func_name_regex = re.compile(r"^(.*?)", re.DOTALL) +_arg_pair_regex = re.compile( + r"(.*?)(?:\\n|\s)*(.*?)", + re.DOTALL, +) -def _strip_names(tool_call): - if isinstance(tool_call, list): - return [_strip_names(tc) for tc in tool_call] - if isinstance(tool_call, dict) and isinstance(tool_call.get("name"), str): - tool_call = dict(tool_call) - tool_call["name"] = tool_call["name"].strip() - return tool_call +def _is_string_type( + tool_name: str, + arg_name: str, + tools: list[Any] | None, +) -> bool: + if tools is None: + return False + for tool in tools: + func = tool.get("function", {}) + if func.get("name") != tool_name: + continue + params = func.get("parameters") or {} + arg_type = params.get("properties", {}).get(arg_name, {}).get("type") + return arg_type == "string" + return False -def parse_tool_call(text: str, tools: list[Any] | None = None): - """Parse Laguna XML-like tool calls into a name and arguments dict.""" - match = glm47._func_name_regex.search(text) + +def _deserialize(value: str) -> Any: + try: + return json.loads(value) + except Exception: + pass + try: + return ast.literal_eval(value) + except Exception: + pass + return value + + +def _parse_single_call(text: str, tools: list[Any] | None): + text = text.strip() + match = _func_name_regex.search(text) if not match: - return _strip_names(glm47.parse_tool_call(text, tools)) + func_name = text.split("\n", 1)[0].strip() + return dict(name=func_name, arguments={}) func_name = match.group(1).strip() - string_args = glm47._get_string_arg_names(func_name, tools) arguments = {} - for match in glm47._func_arg_regex.finditer(text): + for match in _arg_pair_regex.finditer(text): arg_key = match.group(1).strip() arg_val = match.group(2).strip() - if arg_key not in string_args: - arg_val = glm47._deserialize(arg_val) + if not _is_string_type(func_name, arg_key, tools): + arg_val = _deserialize(arg_val) arguments[arg_key] = arg_val return dict(name=func_name, arguments=arguments) + + +def parse_tool_call(text: str, tools: list[Any] | None = None): + matches = _tool_call_regex.findall(text) + if matches: + calls = [_parse_single_call(match, tools) for match in matches] + return calls[0] if len(calls) == 1 else calls + + return _parse_single_call(text, tools) diff --git a/tests/test_models.py b/tests/test_models.py index a9418ca1d..3cc031707 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1859,6 +1859,55 @@ def test_hunyuan(self): model, args.model_type, args.vocab_size, args.num_hidden_layers ) + def test_laguna(self): + from mlx_lm.models import laguna + + args = laguna.ModelArgs( + model_type="laguna", + vocab_size=1000, + hidden_size=128, + intermediate_size=256, + num_hidden_layers=4, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=16, + max_position_embeddings=1000, + rms_norm_eps=1e-5, + attention_bias=False, + mlp_only_layers=[0], + num_experts=4, + num_experts_per_tok=2, + decoder_sparse_step=1, + moe_intermediate_size=64, + shared_expert_intermediate_size=64, + norm_topk_prob=True, + moe_routed_scaling_factor=2.5, + gating="per-head", + sliding_window=4, + layer_types=[ + "full_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + ], + num_attention_heads_per_layer=[4, 8, 8, 4], + rope_parameters={ + "rope_theta": 10000.0, + "rope_type": "default", + "partial_rotary_factor": 0.5, + }, + swa_rope_parameters={ + "rope_theta": 10000.0, + "rope_type": "linear", + "factor": 1.0, + "partial_rotary_factor": 1.0, + }, + ) + model = laguna.Model(args) + self.model_test_runner( + model, args.model_type, args.vocab_size, args.num_hidden_layers + ) + def test_hunyuan_v1_dense(self): from mlx_lm.models import hunyuan_v1_dense @@ -2647,47 +2696,6 @@ def test_all_models(self): "max_position_embeddings": 1000, "norm_topk_prob": True, }, - { - "model_type": "laguna", - "vocab_size": 1000, - "hidden_size": 128, - "intermediate_size": 256, - "num_hidden_layers": 4, - "num_attention_heads": 4, - "num_key_value_heads": 2, - "head_dim": 16, - "max_position_embeddings": 1000, - "rms_norm_eps": 1e-5, - "num_experts": 4, - "num_experts_per_tok": 2, - "decoder_sparse_step": 1, - "mlp_only_layers": [0], - "moe_intermediate_size": 64, - "shared_expert_intermediate_size": 64, - "norm_topk_prob": True, - "moe_routed_scaling_factor": 2.5, - "gating": "per-head", - "sliding_window": 4, - "layer_types": [ - "full_attention", - "sliding_attention", - "sliding_attention", - "full_attention", - ], - "num_attention_heads_per_layer": [4, 8, 8, 4], - "rope_parameters": { - "rope_theta": 10000.0, - "rope_type": "default", - "partial_rotary_factor": 0.5, - }, - "swa_rope_parameters": { - "rope_theta": 10000.0, - "rope_type": "linear", - "factor": 1.0, - "partial_rotary_factor": 1.0, - }, - "tie_word_embeddings": False, - }, { "model_type": "lille-130m", "block_size": 128, diff --git a/tests/test_tool_parsing.py b/tests/test_tool_parsing.py index f134edafe..7eff6fae4 100644 --- a/tests/test_tool_parsing.py +++ b/tests/test_tool_parsing.py @@ -30,6 +30,14 @@ def test_parsers(self): "multiplya12234585b48838483920", glm47, ), + ( + "multiply\n" + "a\n" + "12234585\n" + "b\n" + "48838483920", + laguna, + ), ( '{"name": "multiply", "arguments": {"a": 12234585, "b": 48838483920}}', json_tools, @@ -100,6 +108,12 @@ def test_parsers(self): 'get_current_temperaturelocation"London"', glm47, ), + ( + "get_current_temperature\n" + "location\n" + "London", + laguna, + ), ( '{"name": "get_current_temperature", "arguments": {"location": "London"}}', json_tools, @@ -365,6 +379,38 @@ def test_laguna_strips_function_name_and_preserves_string_args(self): {"location": "Warsaw", "days": 3, "postal_code": "00123"}, ) + def test_laguna_full_tool_call_tags(self): + test_case = ( + "search\n" + "query\n" + "weather\n" + "" + ) + tool_call = laguna.parse_tool_call(test_case, None) + self.assertEqual( + tool_call, + {"name": "search", "arguments": {"query": "weather"}}, + ) + + test_case = ( + "search\n" + "query\n" + "weather\n" + "" + "read_file\n" + "path\n" + "/tmp/test.txt\n" + "" + ) + tool_calls = laguna.parse_tool_call(test_case, None) + self.assertEqual( + tool_calls, + [ + {"name": "search", "arguments": {"query": "weather"}}, + {"name": "read_file", "arguments": {"path": "/tmp/test.txt"}}, + ], + ) + if __name__ == "__main__": unittest.main() From 0857ee1cf1f4ba7c43e73b836309d9c884529ca8 Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Tue, 28 Apr 2026 17:34:01 +0200 Subject: [PATCH 8/8] Refactor ModelArgs and update gating logic --- mlx_lm/models/laguna.py | 54 ++++++++++++++++++++++++++++++++++++----- tests/test_models.py | 22 +++++++++-------- 2 files changed, 60 insertions(+), 16 deletions(-) diff --git a/mlx_lm/models/laguna.py b/mlx_lm/models/laguna.py index 79ca4325d..515ec8230 100644 --- a/mlx_lm/models/laguna.py +++ b/mlx_lm/models/laguna.py @@ -27,14 +27,14 @@ class ModelArgs(BaseModelArgs): gating: Union[bool, str] = True tie_word_embeddings: bool = False rope_theta: float = 500000.0 - rope_parameters: Optional[Dict[str, Union[float, str]]] = None - rope_scaling: Optional[Dict[str, Union[float, str]]] = None + rope_parameters: Optional[Dict[str, Any]] = None + rope_scaling: Optional[Dict[str, Any]] = None partial_rotary_factor: Optional[float] = None rope_style: str = "rotate-half" sliding_window: Optional[int] = None layer_types: Optional[List[str]] = None num_attention_heads_per_layer: Optional[List[int]] = None - swa_rope_parameters: Optional[Dict[str, Union[float, str]]] = None + swa_rope_parameters: Optional[Dict[str, Any]] = None swa_attention_sink_enabled: bool = False num_experts: int = 0 num_experts_per_tok: int = 0 @@ -49,6 +49,9 @@ class ModelArgs(BaseModelArgs): moe_router_use_sigmoid: bool = True def __post_init__(self): + if self.gating is True: + self.gating = "per-head" + if self.layer_types is None: self.layer_types = ["full_attention"] * self.num_hidden_layers if len(self.layer_types) != self.num_hidden_layers: @@ -69,18 +72,57 @@ def __post_init__(self): "Every query-head count must be divisible by num_key_value_heads." ) - if self.rope_parameters is None: - self.rope_parameters = ( + rope_parameters = ( + dict(self.rope_parameters) + if self.rope_parameters is not None + else ( dict(self.rope_scaling) if self.rope_scaling is not None else {"rope_type": "default", "rope_theta": self.rope_theta} ) + ) + + layer_types = set(self.layer_types) + layer_rope_parameters = { + k: v + for k, v in rope_parameters.items() + if k in layer_types and isinstance(v, dict) + } + if layer_rope_parameters: + top_level_parameters = { + k: v + for k, v in rope_parameters.items() + if k not in layer_types and not isinstance(v, dict) + } + + def rope_parameters_for(layer_type: str) -> Dict[str, Any]: + params = dict(layer_rope_parameters.get(layer_type, {})) + for k, v in top_level_parameters.items(): + params.setdefault(k, v) + return params + + default_layer_type = ( + "full_attention" + if "full_attention" in layer_rope_parameters + else next(iter(layer_rope_parameters)) + ) + self.rope_parameters = rope_parameters_for(default_layer_type) + + if ( + self.swa_rope_parameters is None + and "sliding_attention" in layer_rope_parameters + ): + self.swa_rope_parameters = rope_parameters_for("sliding_attention") else: - self.rope_parameters = dict(self.rope_parameters) + self.rope_parameters = rope_parameters if self.swa_rope_parameters is not None: self.swa_rope_parameters = dict(self.swa_rope_parameters) + self.rope_parameters.setdefault("rope_type", "default") + if self.swa_rope_parameters is not None: + self.swa_rope_parameters.setdefault("rope_type", "default") + if self.partial_rotary_factor is not None: self.rope_parameters.setdefault( "partial_rotary_factor", self.partial_rotary_factor diff --git a/tests/test_models.py b/tests/test_models.py index 3cc031707..a7cc96a4d 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1882,7 +1882,7 @@ def test_laguna(self): shared_expert_intermediate_size=64, norm_topk_prob=True, moe_routed_scaling_factor=2.5, - gating="per-head", + gating=True, sliding_window=4, layer_types=[ "full_attention", @@ -1892,15 +1892,17 @@ def test_laguna(self): ], num_attention_heads_per_layer=[4, 8, 8, 4], rope_parameters={ - "rope_theta": 10000.0, - "rope_type": "default", - "partial_rotary_factor": 0.5, - }, - swa_rope_parameters={ - "rope_theta": 10000.0, - "rope_type": "linear", - "factor": 1.0, - "partial_rotary_factor": 1.0, + "full_attention": { + "rope_theta": 10000.0, + "rope_type": "default", + "partial_rotary_factor": 0.5, + }, + "sliding_attention": { + "rope_theta": 10000.0, + "rope_type": "linear", + "factor": 1.0, + "partial_rotary_factor": 1.0, + }, }, ) model = laguna.Model(args)