@modular-prompt/driver 0.16.0 → 0.17.0

package/src/mlx-ml/python/handlers/cache.py

@@ -18,6 +18,8 @@ def handle_cache_prefill(
     prefix_hashes: list[str] | None = None,
     tools: list | None = None,
     reasoning_effort: str | None = None,
+    images: list | None = None,
+    max_image_size: int = 768,
 ) -> None:
     tokenizer = backend.get_tokenizer()
 
@@ -74,6 +76,8 @@ def handle_cache_prefill(
         trim_to_tokens=trim_to_tokens,
         prefix_offsets=prefix_offsets,
         prefix_hashes=prefix_hashes,
+        images=images,
+        max_image_size=max_image_size,
     )
     if prefix_offsets and prefix_hashes:
         result["prefix_offsets"] = prefix_offsets

package/src/mlx-ml/python/handlers/generate.py

@@ -6,7 +6,6 @@ import re
 import sys
 
 from backends.base import ModelBackend
-from mlx_lm.models.cache import trim_prompt_cache
 from handlers.cancel import poll_cancel
 
 
@@ -29,6 +28,7 @@ def _stream_to_stdout(
     images: list | None = None,
     primer: str | None = None,
     prompt_cache: list | None = None,
+    cache_loaded: bool | None = None,
 ) -> None:
     if primer is not None:
         print(primer, end="", flush=True)
@@ -46,6 +46,8 @@ def _stream_to_stdout(
             meta["prompt_tokens"] = last_response.prompt_tokens
         if hasattr(last_response, "generation_tokens"):
             meta["generation_tokens"] = last_response.generation_tokens
+    if cache_loaded is not None:
+        meta["cache_loaded"] = cache_loaded
 
     if meta:
         print(f"\x1e__META__:{json.dumps(meta)}", end="\0", flush=True)
@@ -84,16 +86,22 @@ def handle_generate(
         else:
             sys.stderr.write(f"--- prompt\n{prompt}\n")
 
-    # VLM 経路では KV キャッシュを使わない（旧 chat ハンドラと同様）
     prompt_cache = None
     cache_tokens = 0
-    if cache_path and not images:
-        prompt_cache = backend.load_cache_from_file(cache_path)
+    cache_loaded = None
+    if cache_path:
+        prompt_cache = backend.load_cache_from_file(
+            cache_path,
+            images=images,
+            max_image_size=max_image_size,
+            prompt=prompt,
+        )
+        cache_loaded = prompt_cache is not None
     if prompt_cache is not None:
         if cache_trim_tokens is not None:
             current_offset = backend.get_cache_offset(prompt_cache)
             if current_offset > cache_trim_tokens:
-                trim_prompt_cache(prompt_cache, current_offset - cache_trim_tokens)
+                backend.trim_cache(prompt_cache, current_offset - cache_trim_tokens)
                 sys.stderr.write(
                     f"KV cache trimmed: {current_offset} → {cache_trim_tokens} tokens\n"
                 )
@@ -101,9 +109,21 @@ def handle_generate(
             else:
                 cache_tokens = current_offset
         else:
-            meta_count = _read_cache_token_count(cache_path) if cache_path else None
+            meta_count = (
+                _read_cache_token_count(cache_path)
+                if cache_path
+                else None
+            )
             if meta_count is not None:
                 cache_tokens = meta_count
+            elif (
+                backend.model_kind == "vlm"
+                and cache_path is not None
+                and cache_path.startswith("mlx-vlm-memory://")
+            ):
+                # Keep the Phase 1 in-process ref as a compatibility fallback.
+                # Persistent VLM refs always have the sidecar above.
+                cache_tokens = backend.get_cache_offset(prompt_cache)
             else:
                 sys.stderr.write(
                     f"WARNING: Cache file exists but no .meta.json found at {cache_path}. "
@@ -111,6 +131,7 @@ def handle_generate(
                 )
                 prompt_cache = None
                 cache_tokens = 0
+                cache_loaded = False
         if prompt_cache is not None:
             sys.stderr.write(
                 f"KV cache loaded: {len(prompt_cache)} layers, {cache_tokens} cached tokens\n"
@@ -122,11 +143,11 @@ def handle_generate(
 
     effective_prompt = prompt
     if prompt_cache is not None and cache_tokens > 0 and isinstance(prompt, str):
-        tokenizer = backend.get_tokenizer()
-        add_special = tokenizer.bos_token is None or not prompt.startswith(
-            tokenizer.bos_token
+        full_tokens = backend.tokenize_prompt(
+            prompt,
+            images=images,
+            max_image_size=max_image_size,
         )
-        full_tokens = tokenizer.encode(prompt, add_special_tokens=add_special)
 
         if cache_tokens < len(full_tokens):
             effective_prompt = full_tokens[cache_tokens:]
@@ -140,6 +161,7 @@ def handle_generate(
                 f"ignoring prefill state\n"
             )
             prompt_cache = None
+            cache_loaded = False
 
     _stream_to_stdout(
         backend,
@@ -148,4 +170,5 @@ def handle_generate(
         images=images,
         primer=primer,
         prompt_cache=prompt_cache,
+        cache_loaded=cache_loaded,
     )

package/src/mlx-ml/python/handlers/tokenize.py

@@ -50,10 +50,7 @@ def handle_tokenize(
             prompt = generate_merged_prompt(messages, capabilities)
 
         # トークン化
-        add_special = tokenizer.bos_token is None or not prompt.startswith(
-            tokenizer.bos_token or ""
-        )
-        token_ids = tokenizer.encode(prompt, add_special_tokens=add_special)
+        token_ids = backend.tokenize_prompt(prompt)
 
         result["token_ids"] = token_ids
         result["token_count"] = len(token_ids)

package/src/mlx-ml/python/pyproject.toml

@@ -10,7 +10,7 @@ dependencies = [
     "mlx>=0.32.2; sys_platform == 'darwin'",
     "mlx-lm==0.31.3; sys_platform == 'darwin'",
     "mlx-optiq>=0.4.4; sys_platform == 'darwin'",
-    "mlx-vlm==0.6.17; sys_platform == 'darwin'",
+    "mlx-vlm==0.7.0; sys_platform == 'darwin'",
     "tokenizers==0.22.2",
     "torch==2.9.1",
     "torchvision==0.24.1",

package/src/mlx-ml/python/server.py

@@ -100,6 +100,8 @@ class Server:
                     prefix_hashes=req.get('prefix_hashes'),
                     tools=req.get('tools'),
                     reasoning_effort=req.get('reasoning_effort'),
+                    images=req.get('images'),
+                    max_image_size=req.get('maxImageSize', 768),
                 )
 
             elif method == 'render':

package/src/mlx-ml/python/uv.lock

@@ -724,22 +724,21 @@ wheels = [
 
 [[package]]
 name = "mlx-audio"
-version = "0.4.3"
+version = "0.5.3"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "huggingface-hub", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" },
     { name = "miniaudio", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" },
     { name = "mlx", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" },
-    { name = "mlx-lm", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" },
     { name = "numpy", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" },
     { name = "scipy", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" },
     { name = "sounddevice", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" },
     { name = "tqdm", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" },
     { name = "transformers", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/95/db/a9f95e3794eca373d681220c8b9f8f84451a0d14959f85cc341ca592394c/mlx_audio-0.4.3.tar.gz", hash = "sha256:8e87badf56a0f73bf91e3797b1195c01440a181cf0b64a2a08dc1bda4b037f54", size = 1144947, upload-time = "2026-04-28T20:18:12.09Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/60/e8/54d02c8bb50fc5b2cae0fbcb459acf78f1fbb316f20f758a8f56824aa407/mlx_audio-0.5.3.tar.gz", hash = "sha256:9e96f1ee68bc81e7ba74d85b8eab213384c8cd8c6f66ccd12034ded7cb5cd167", size = 1709447, upload-time = "2026-09-07T19:44:40.285Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/e5/25/0a89073ed7b7cdf34299042bd03d867c12c0c8b43f597be61bea7f146793/mlx_audio-0.4.3-py3-none-any.whl", hash = "sha256:6b87bf42d79d9ceb6b9310a77656b9b76429c2d6ddd89f634b2786c58a2e4721", size = 1373582, upload-time = "2026-04-28T20:18:10.512Z" },
+    { url = "https://files.pythonhosted.org/packages/84/c2/2876e6ff4811b1d6a43e1480947dd6580d02dba16f34724b2e3645090997/mlx_audio-0.5.3-py3-none-any.whl", hash = "sha256:8d920b2dcbcf37b5fd2bbb070e21a9a33442bf231a754e0d39387e65986bd830", size = 2057098, upload-time = "2026-09-07T19:44:38.581Z" },
 ]
 
 [[package]]
@@ -773,7 +772,7 @@ requires-dist = [
     { name = "mlx", marker = "sys_platform == 'darwin'", specifier = ">=0.32.2" },
     { name = "mlx-lm", marker = "sys_platform == 'darwin'", specifier = "==0.31.3" },
     { name = "mlx-optiq", marker = "sys_platform == 'darwin'", specifier = ">=0.4.4" },
-    { name = "mlx-vlm", marker = "sys_platform == 'darwin'", specifier = "==0.6.17" },
+    { name = "mlx-vlm", marker = "sys_platform == 'darwin'", specifier = "==0.7.0" },
     { name = "tokenizers", specifier = "==0.22.2" },
     { name = "torch", specifier = "==2.9.1" },
     { name = "torchvision", specifier = "==0.24.1" },
@@ -846,10 +845,11 @@ wheels = [
 
 [[package]]
 name = "mlx-vlm"
-version = "0.6.17"
+version = "0.7.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "fastapi", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" },
+    { name = "jinja2", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" },
     { name = "llguidance", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" },
     { name = "miniaudio", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" },
     { name = "mlx", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" },
@@ -866,9 +866,9 @@ dependencies = [
     { name = "uvicorn", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" },
     { name = "websockets", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/c5/a1/4b21ea85918e5223144bbe80cb649f5953c105ea440f0eba45dd00a5b863/mlx_vlm-0.6.17.tar.gz", hash = "sha256:5e4ff5426af093b3cf41b4e5e33c272db82fa33000cb3c36a1aab9e218ad7ee7", size = 2043597, upload-time = "2026-08-26T17:08:38.01Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/1e/9e/220a9427b1cf6f4ee381aa475fc47ecbe6d84f4177f0ac28176c3435fbdd/mlx_vlm-0.7.0.tar.gz", hash = "sha256:0bfc3ae1f82cb0a7866bf9f1257d672f6450c5fbd13f22b2d62fddb72b759379", size = 2208403, upload-time = "2026-09-07T15:25:14.497Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/f8/42/f275566dba66b54c21e039344c4fb90fd3a3749fc251754545249b45dc8d/mlx_vlm-0.6.17-py3-none-any.whl", hash = "sha256:d24723244ef89b2a6140bec45e38c078689ea89450edfc0e348ad149f5af3b1c", size = 2590205, upload-time = "2026-08-26T17:08:36.148Z" },
+    { url = "https://files.pythonhosted.org/packages/65/2c/3868b7ccb50c3cb975e7387de0b1e642b8896f2504a80d097f109d694eb4/mlx_vlm-0.7.0-py3-none-any.whl", hash = "sha256:5ea0c2b8182c055068cb0da59cb2503c238c428a07daaa0bebc246e8c666f4d3", size = 2787810, upload-time = "2026-09-07T15:25:12.683Z" },
 ]
 
 [[package]]

package/src/pytorch/templates/cpu-minimal/backends/base.py

@@ -0,0 +1,139 @@
+from abc import ABC, abstractmethod
+from typing import Any, Iterator
+
+
+class ModelBackend(ABC):
+    """Abstract base class for model backends."""
+
+    @abstractmethod
+    def load(self, model_name: str) -> None:
+        """Load the target model."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_tokenizer(self) -> Any:
+        """Return the tokenizer or processor."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def stream_generate(
+        self, prompt: str | list[int], options: dict, images: list | None = None,
+        prompt_cache: Any | None = None,
+    ) -> Iterator[Any]:
+        """Stream generation results."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def supports_vision(self) -> bool:
+        """Return whether image input is supported."""
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def model_kind(self) -> str:
+        """Return "lm" or "vlm"."""
+        raise NotImplementedError
+
+    def load_drafter(self, drafter_model: str) -> None:
+        """Load a drafter model for speculative decoding."""
+        raise NotImplementedError(
+            f"{type(self).__name__} does not support drafter models"
+        )
+
+    def has_drafter(self) -> bool:
+        """Return whether a drafter model is loaded."""
+        return False
+
+    def cache_prefill(
+        self,
+        cache_path: str,
+        prompt: str,
+        base_cache_path: str | None = None,
+        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 consume_cache_write_tokens(self, cache_path: str) -> int:
+        """Return and clear write usage pending for a cache reference."""
+        return 0
+
+    def trim_cache(self, prompt_cache: Any, tokens: int) -> Any:
+        """Remove trailing tokens from a backend-owned prompt cache."""
+        raise NotImplementedError(
+            f"{type(self).__name__} does not support prompt cache trimming"
+        )
+
+    def tokenize_prompt(
+        self,
+        prompt: str,
+        images: list | None = None,
+        max_image_size: int = 768,
+    ) -> list[int]:
+        """Tokenize a rendered prompt using the backend's prompt rules."""
+        tokenizer = self.get_tokenizer()
+        bos_token = getattr(tokenizer, "bos_token", None)
+        add_special = bos_token is None or not prompt.startswith(bos_token or "")
+        token_ids = tokenizer.encode(prompt, add_special_tokens=add_special)
+        if hasattr(token_ids, "flatten"):
+            token_ids = token_ids.flatten().tolist()
+        return [int(token_id) for token_id in token_ids]
+
+    def load_cache_from_file(
+        self,
+        cache_path: str,
+        images: list | None = None,
+        max_image_size: int = 768,
+        prompt: str | list[int] | None = None,
+        prefix_token_count: int | None = None,
+    ) -> Any | None:
+        """Load a prompt cache, optionally validating only a prompt prefix."""
+        return None
+
+    def get_cache_offset(self, prompt_cache: Any) -> int:
+        """Get the number of tokens stored in a loaded prompt cache."""
+        if not prompt_cache:
+            return 0
+
+        get_seq_length = getattr(prompt_cache, "get_seq_length", None)
+        if callable(get_seq_length):
+            try:
+                return int(get_seq_length())
+            except Exception:
+                pass
+
+        keys = getattr(prompt_cache, "keys", None)
+        if keys is not None:
+            try:
+                return int(keys.shape[-2])
+            except Exception:
+                pass
+
+        layers = getattr(prompt_cache, "layers", None)
+        if layers:
+            return self.get_cache_offset(layers[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:
+            key = layer0[0] if isinstance(layer0, (list, tuple)) else layer0
+            shape = key.shape
+            return int(shape[-2] if len(shape) >= 2 else shape[0])
+        except Exception:
+            pass
+        if hasattr(layer0, 'keys') and layer0.keys is not None:
+            return int(layer0.keys.shape[-2])
+        return 0