From 52fafa7396655770c5bf06c307ba5f53df067914 Mon Sep 17 00:00:00 2001 From: Pierre Lamy Date: Mon, 6 Jul 2026 14:40:33 -0400 Subject: [PATCH] laguna: windowed KV via make_cache() + strip language_model. prefix Two small additions on top of the Laguna model: - make_cache(): the sliding_attention layers (the majority) never attend beyond their window, so return a bounded RotatingKVCache(max_size= sliding_window) for them and a full KVCache only for the full_attention (global) layers. Without this, make_prompt_cache allocates a full cache on every layer, so long-context runs store and attend over history the sliding layers can't use (measured ~2.7x decode and ~4 GB less KV at 32k on the 8-bit repack); bitwise-identical to a full cache at <= window context. - sanitize(): strip the VLM-style `language_model.` prefix that some repacks (e.g. AtomicChat/Laguna-XS-2.1-MLX-8bit) put on every tensor, so a stock load finds the weights. Co-Authored-By: Claude Opus 4.8 --- mlx_lm/models/laguna.py | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/mlx_lm/models/laguna.py b/mlx_lm/models/laguna.py index 515ec8230..5906bd69d 100644 --- a/mlx_lm/models/laguna.py +++ b/mlx_lm/models/laguna.py @@ -6,6 +6,7 @@ from .activations import swiglu from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention +from .cache import KVCache, RotatingKVCache from .rope_utils import initialize_rope from .switch_layers import SwitchGLU @@ -406,7 +407,33 @@ def __call__( return self.model.embed_tokens.as_linear(out) return self.lm_head(out) + def make_cache(self): + # Most Laguna layers are sliding_attention: they never attend beyond + # their window, so a bounded RotatingKVCache is both correct and far + # cheaper than a full cache at long context. Only the full_attention + # (global) layers need an unbounded KVCache. + caches = [] + for lt in self.args.layer_types: + # A sliding layer needs a valid window; if a (malformed) config + # omits sliding_window, fall back to a full cache rather than build + # a RotatingKVCache with max_size=None. + if lt == "sliding_attention" and self.args.sliding_window: + caches.append(RotatingKVCache(max_size=self.args.sliding_window)) + else: + caches.append(KVCache()) + return caches + def sanitize(self, weights): + # Some repacks (e.g. AtomicChat/Laguna-XS-2.1-MLX-8bit) wrap every + # tensor under a VLM-style `language_model.` prefix. Strip it so the + # keys line up with this module tree (model.* / lm_head.*). + if any(k.startswith("language_model.") for k in weights): + prefix = "language_model." + weights = { + (k[len(prefix) :] if k.startswith(prefix) else k): v + for k, v in weights.items() + } + if self.args.tie_word_embeddings: weights.pop("lm_head.weight", None)