@modular-prompt/driver 0.15.0 → 0.17.0

package/src/mlx-ml/python/backends/base.py

@@ -52,33 +52,103 @@ class ModelBackend(ABC):
         trim_to_tokens: int | None = None,
         prefix_offsets: list[int] | None = None,
         prefix_hashes: list[str] | None = None,
+        images: list | None = None,
+        max_image_size: int = 768,
     ) -> dict:
         """Build a KV cache from a prompt prefix."""
         raise NotImplementedError(
             f"{type(self).__name__} does not support prompt caching"
         )
 
-    def load_cache_from_file(self, cache_path: str) -> list | None:
-        """Load a prompt cache from file, or None."""
+    def tokenize_prompt(
+        self,
+        prompt: str,
+        images: list | None = None,
+        max_image_size: int = 768,
+    ) -> list[int]:
+        """Tokenize a rendered prompt using this backend's prompt rules.
+
+        A VLM backend returns a processor from ``get_tokenizer()``.  Keeping
+        this operation on the backend prevents shared handlers from assuming
+        that the returned object is a plain tokenizer.  ``images`` and
+        ``max_image_size`` let multimodal backends reproduce the same image
+        placeholder expansion as their generation path; text-only backends
+        ignore them.
+        """
+        tokenizer = self.get_tokenizer()
+        tokenizer = getattr(tokenizer, "tokenizer", tokenizer)
+        bos_token = getattr(tokenizer, "bos_token", None)
+        add_special = bos_token is None or not prompt.startswith(bos_token or "")
+        return list(tokenizer.encode(prompt, add_special_tokens=add_special))
+
+    def trim_cache(self, prompt_cache: list, tokens: int) -> None:
+        """Trim a backend-owned prompt cache in place.
+
+        Cache implementations are backend-specific.  The default handles the
+        cache objects exposed by mlx-vlm without importing mlx-lm into the
+        shared generate handler.
+        """
+        if tokens <= 0:
+            return
+
+        for entry in prompt_cache:
+            trim = getattr(entry, "trim", None)
+            if callable(trim):
+                trim(tokens)
+                continue
+
+            children = getattr(entry, "caches", None)
+            if children is not None:
+                for child in children:
+                    child_trim = getattr(child, "trim", None)
+                    if callable(child_trim):
+                        child_trim(tokens)
+
+    def load_cache_from_file(
+        self,
+        cache_path: str,
+        images: list | None = None,
+        max_image_size: int = 768,
+        prompt: str | list[int] | None = None,
+    ) -> list | None:
+        """Load a prompt cache from file, or None.
+
+        ``prompt`` is optional so backends that persist prompt identities can
+        validate the loaded cache against the current request.
+        """
         return None
 
     def get_cache_offset(self, prompt_cache: list) -> int:
         """Get the number of tokens stored in a loaded prompt cache."""
         if not prompt_cache:
             return 0
-        layer0 = prompt_cache[0]
-        if hasattr(layer0, 'offset'):
-            off = layer0.offset
-            return int(off.item() if hasattr(off, 'item') else off)
-        if hasattr(layer0, 'caches'):
-            for c in layer0.caches:
-                if hasattr(c, 'offset'):
-                    off = c.offset
-                    return int(off.item() if hasattr(off, 'item') else off)
-        try:
-            return int(layer0[0].shape[2])
-        except Exception:
-            pass
-        if hasattr(layer0, 'keys') and layer0.keys is not None:
-            return int(layer0.keys.shape[2])
-        return 0
+
+        def offset_of(cache: Any) -> int:
+            if hasattr(cache, "offset"):
+                off = cache.offset
+                return int(off.item() if hasattr(off, "item") else off)
+
+            size = getattr(cache, "size", None)
+            if callable(size):
+                try:
+                    return int(size())
+                except Exception:
+                    pass
+
+            children = getattr(cache, "caches", None)
+            if children is not None:
+                return max((offset_of(child) for child in children), default=0)
+
+            keys = getattr(cache, "keys", None)
+            if keys is not None:
+                try:
+                    return int(keys.shape[-2])
+                except Exception:
+                    pass
+
+            try:
+                return int(cache[0].shape[2])
+            except Exception:
+                return 0
+
+        return max((offset_of(layer) for layer in prompt_cache), default=0)

package/src/mlx-ml/python/backends/cache_archive.py

@@ -0,0 +1,41 @@
+"""Streaming zip storage for MLX prompt caches."""
+
+from __future__ import annotations
+
+import zipfile
+from collections.abc import Callable
+from typing import Any
+
+
+CACHE_ENTRY_NAME = "prompt_cache.safetensors"
+
+
+def save_prompt_cache_zip(
+    file_name: str,
+    cache: Any,
+    save_impl: Callable[[Any, Any], None],
+) -> None:
+    """Save a prompt cache into a compressed zip member as it is produced.
+
+    ``save_impl`` receives the zip member's writable stream.  It must write
+    the safetensors payload to that stream instead of creating an intermediate
+    uncompressed file.
+    """
+    with zipfile.ZipFile(
+        file_name,
+        mode="w",
+        compression=zipfile.ZIP_DEFLATED,
+        allowZip64=True,
+    ) as archive:
+        with archive.open(CACHE_ENTRY_NAME, mode="w", force_zip64=True) as cache_file:
+            save_impl(cache_file, cache)
+
+
+def load_prompt_cache_zip(
+    file_name: str,
+    load_impl: Callable[[Any], Any],
+) -> Any:
+    """Load a prompt cache from the safetensors member of a zip archive."""
+    with zipfile.ZipFile(file_name, mode="r") as archive:
+        with archive.open(CACHE_ENTRY_NAME, mode="r") as cache_file:
+            return load_impl(cache_file)

package/src/mlx-ml/python/backends/mlx_lm.py

@@ -7,13 +7,29 @@ from typing import Any, Iterator
 
 from mlx_lm import load as mlx_lm_load
 from mlx_lm import stream_generate as mlx_lm_stream_generate
-from mlx_lm.models.cache import make_prompt_cache, save_prompt_cache, load_prompt_cache, trim_prompt_cache
+from mlx_lm.models.cache import (
+    make_prompt_cache,
+    save_prompt_cache as mlx_save_prompt_cache,
+    load_prompt_cache as mlx_load_prompt_cache,
+    trim_prompt_cache,
+)
 from mlx_lm.sample_utils import make_sampler
 
 from backends.base import ModelBackend
+from backends.cache_archive import load_prompt_cache_zip, save_prompt_cache_zip
 from utils.token_utils import is_eod_token
 
 
+def save_prompt_cache(file_name: str, cache: Any) -> None:
+    """Save a prompt cache as a compressed archive without an uncompressed copy."""
+    save_prompt_cache_zip(file_name, cache, mlx_save_prompt_cache)
+
+
+def load_prompt_cache(file_name: str) -> list:
+    """Load a prompt cache from its compressed safetensors archive."""
+    return load_prompt_cache_zip(file_name, mlx_load_prompt_cache)
+
+
 class MlxLmBackend(ModelBackend):
     """`mlx_lm` backend for text-only models."""
 
@@ -62,12 +78,29 @@ class MlxLmBackend(ModelBackend):
 
     # get_cache_offset is inherited from ModelBackend base class
 
-    def _tokenize_prompt(self, prompt: str) -> list[int]:
+    def _tokenize_prompt(
+        self,
+        prompt: str,
+        images: list | None = None,
+        max_image_size: int = 768,
+    ) -> list[int]:
         """Tokenize a prompt string using the same logic as stream_generate."""
         add_special = self.tokenizer.bos_token is None or not prompt.startswith(
             self.tokenizer.bos_token
         )
-        return self.tokenizer.encode(prompt, add_special_tokens=add_special)
+        return list(self.tokenizer.encode(prompt, add_special_tokens=add_special))
+
+    def tokenize_prompt(
+        self,
+        prompt: str,
+        images: list | None = None,
+        max_image_size: int = 768,
+    ) -> list[int]:
+        return self._tokenize_prompt(prompt, images, max_image_size)
+
+    def trim_cache(self, prompt_cache: list, tokens: int) -> None:
+        if tokens > 0:
+            trim_prompt_cache(prompt_cache, tokens)
 
     @staticmethod
     def _write_cache_meta(
@@ -106,6 +139,8 @@ class MlxLmBackend(ModelBackend):
         trim_to_tokens: int | None = None,
         prefix_offsets: list[int] | None = None,
         prefix_hashes: list[str] | None = None,
+        images: list | None = None,
+        max_image_size: int = 768,
     ) -> dict:
         if self.model is None or self.tokenizer is None:
             raise RuntimeError("Model is not loaded")
@@ -184,7 +219,13 @@ class MlxLmBackend(ModelBackend):
             sys.stderr.write(f"Cache created: {cache_path} ({token_count} tokens)\n")
         return {"cache_path": cache_path, "token_count": token_count}
 
-    def load_cache_from_file(self, cache_path: str) -> list | None:
+    def load_cache_from_file(
+        self,
+        cache_path: str,
+        images: list | None = None,
+        max_image_size: int = 768,
+        prompt: str | list[int] | None = None,
+    ) -> list | None:
         try:
             return load_prompt_cache(cache_path)
         except FileNotFoundError: