@modular-prompt/driver 0.16.0 → 0.17.1

package/src/pytorch/{python → templates/cpu-minimal}/handlers/__init__.py

@@ -1,3 +1,4 @@
+from handlers.cache import handle_cache_prefill
 from handlers.capabilities import handle_capabilities
 from handlers.completion import handle_completion
 from handlers.format_test import handle_format_test

package/src/pytorch/templates/cpu-minimal/handlers/cache.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+import json
+
+from backends.base import ModelBackend
+from utils.prompt_builder import generate_merged_prompt, supports_chat_template
+
+
+def _render_prefill_prompt(
+    backend: ModelBackend,
+    capabilities: dict,
+    messages: list,
+    tools: list | None,
+    reasoning_effort: str | None,
+) -> str:
+    tokenizer = backend.get_tokenizer()
+    extra_kwargs = {}
+    if tools is not None:
+        extra_kwargs["tools"] = tools
+    if reasoning_effort is not None:
+        extra_kwargs["reasoning_effort"] = reasoning_effort
+
+    if not supports_chat_template(tokenizer):
+        return generate_merged_prompt(messages, capabilities)
+
+    try:
+        return tokenizer.apply_chat_template(
+            messages,
+            add_generation_prompt=False,
+            tokenize=False,
+            **extra_kwargs,
+        )
+    except TypeError:
+        try:
+            fallback_kwargs = {"tools": tools} if tools is not None else {}
+            return tokenizer.apply_chat_template(
+                messages,
+                add_generation_prompt=False,
+                tokenize=False,
+                **fallback_kwargs,
+            )
+        except TypeError:
+            return tokenizer.apply_chat_template(
+                messages,
+                add_generation_prompt=False,
+                tokenize=False,
+            )
+
+
+def handle_cache_prefill(
+    backend: ModelBackend,
+    capabilities: dict,
+    cache_path: str,
+    messages: list,
+    base_cache_path: str | None = None,
+    trim_to_tokens: int | None = None,
+    prefix_offsets: list[int] | None = None,
+    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:
+    """Build a persistent or process-local PyTorch KV cache from chat messages."""
+    if images:
+        raise ValueError("PyTorch LIP backend does not support vision input")
+
+    prompt = _render_prefill_prompt(
+        backend,
+        capabilities,
+        messages,
+        tools,
+        reasoning_effort,
+    )
+    result = backend.cache_prefill(
+        cache_path,
+        prompt,
+        base_cache_path=base_cache_path,
+        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 is not None and prefix_hashes is not None:
+        result["prefix_offsets"] = prefix_offsets
+        result["prefix_hashes"] = prefix_hashes
+    print(json.dumps(result), end="\0", flush=True)

package/src/pytorch/templates/cpu-minimal/handlers/generate.py

@@ -0,0 +1,157 @@
+from __future__ import annotations
+
+import json
+
+from backends.base import ModelBackend
+from handlers.cancel import poll_cancel
+
+
+def _stream_to_stdout(
+    backend: ModelBackend,
+    prompt: str | list[int],
+    options: dict,
+    images: list | None = None,
+    primer: str | None = None,
+    prompt_cache=None,
+    cache_loaded: bool | None = None,
+    cache_read_tokens: int = 0,
+    cache_write_tokens: int = 0,
+) -> None:
+    if images:
+        raise ValueError("PyTorch LIP backend does not support vision input")
+
+    if primer is not None:
+        print(primer, end="", flush=True)
+
+    first_prompt_tokens = None
+    first_cache_read_tokens = None
+    first_cache_write_tokens = None
+    reported_generation_tokens = None
+    for response in backend.stream_generate(
+        prompt,
+        options,
+        images,
+        prompt_cache=prompt_cache,
+    ):
+        if poll_cancel():
+            break
+        response_prompt_tokens = getattr(response, "prompt_tokens", None)
+        if first_prompt_tokens is None and response_prompt_tokens is not None:
+            first_prompt_tokens = response_prompt_tokens
+        if (
+            first_cache_read_tokens is None
+            and getattr(response, "cache_read_tokens", None) is not None
+        ):
+            first_cache_read_tokens = response.cache_read_tokens
+        if (
+            first_cache_write_tokens is None
+            and getattr(response, "cache_write_tokens", None) is not None
+        ):
+            first_cache_write_tokens = response.cache_write_tokens
+        response_generation_tokens = getattr(response, "generation_tokens", None)
+        if response_generation_tokens is not None:
+            reported_generation_tokens = max(
+                reported_generation_tokens or 0,
+                int(response_generation_tokens),
+            )
+        print(response.text.replace("\0", "").replace("\x1e", ""), end="", flush=True)
+
+    meta: dict = {}
+    if first_prompt_tokens is not None:
+        meta["prompt_tokens"] = first_prompt_tokens
+    if reported_generation_tokens is not None:
+        meta["generation_tokens"] = max(0, reported_generation_tokens)
+    if first_cache_read_tokens is not None:
+        meta["cache_read_tokens"] = first_cache_read_tokens
+    if first_cache_write_tokens is not None:
+        meta["cache_write_tokens"] = first_cache_write_tokens
+    if cache_read_tokens > 0 and "cache_read_tokens" not in meta:
+        meta["cache_read_tokens"] = cache_read_tokens
+    if cache_write_tokens > 0 and "cache_write_tokens" not in meta:
+        meta["cache_write_tokens"] = cache_write_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)
+    else:
+        print("", end="\0", flush=True)
+
+
+def handle_generate(
+    backend: ModelBackend,
+    prompt: str | list[int],
+    options: dict | None = None,
+    images: list | None = None,
+    max_image_size: int = 768,
+    primer: str | None = None,
+    cache_path: str | None = None,
+    cache_trim_tokens: int | None = None,
+) -> None:
+    """LIP generate: 整形済み prompt のストリーム推論"""
+    if cache_trim_tokens is not None and cache_trim_tokens < 0:
+        raise ValueError("cache_trim_tokens must be non-negative")
+
+    if options is None:
+        options = {}
+
+    final_options = dict(options)
+    final_options.pop("trust_remote_code", None)
+
+    prompt_cache = None
+    cache_loaded = None
+    cache_read_tokens = 0
+    cache_write_tokens = 0
+    if cache_path:
+        if images:
+            cache_loaded = False
+        else:
+            prompt_cache = backend.load_cache_from_file(
+                cache_path,
+                images=images,
+                max_image_size=max_image_size,
+                prompt=prompt,
+                prefix_token_count=cache_trim_tokens,
+            )
+            cache_loaded = prompt_cache is not None
+
+        if prompt_cache is not None:
+            cache_read_tokens = backend.get_cache_offset(prompt_cache)
+            if cache_trim_tokens is not None and cache_read_tokens > cache_trim_tokens:
+                prompt_cache = backend.trim_cache(
+                    prompt_cache,
+                    cache_read_tokens - cache_trim_tokens,
+                )
+                cache_read_tokens = cache_trim_tokens
+            if cache_read_tokens <= 0:
+                prompt_cache = None
+                cache_loaded = False
+            elif isinstance(prompt, (str, list)):
+                full_tokens = (
+                    backend.tokenize_prompt(prompt)
+                    if isinstance(prompt, str)
+                    else [int(token_id) for token_id in prompt]
+                )
+                if cache_read_tokens < len(full_tokens):
+                    prompt = full_tokens[cache_read_tokens:]
+                else:
+                    # A cache covering the complete prompt cannot be passed
+                    # with an empty input_ids tensor.  The safe fallback is a
+                    # cold generation for this handler.
+                    prompt_cache = None
+                    cache_loaded = False
+                    cache_read_tokens = 0
+                if prompt_cache is not None:
+                    cache_write_tokens = backend.consume_cache_write_tokens(cache_path)
+
+    _stream_to_stdout(
+        backend,
+        prompt,
+        final_options,
+        images=images,
+        primer=primer,
+        prompt_cache=prompt_cache,
+        cache_loaded=cache_loaded,
+        cache_read_tokens=cache_read_tokens,
+        cache_write_tokens=cache_write_tokens,
+    )

package/src/pytorch/{python → templates/cpu-minimal}/pyproject.toml

@@ -4,9 +4,9 @@ version = "0.1.0"
 description = "PyTorch (Transformers) driver for modular-prompt — cpu-minimal runtime"
 requires-python = ">=3.10,<3.14"
 dependencies = [
-    "safetensors==0.7.0",
+    "safetensors==0.8.0",
     "tokenizers==0.22.2",
-    "transformers==4.57.6",
+    "transformers>=5.14.0",
 ]
 
 [dependency-groups]

package/src/pytorch/{python → templates/cpu-minimal}/server.py

@@ -3,7 +3,7 @@ import json
 import sys
 
 from backends.base import ModelBackend
-from handlers import handle_capabilities, handle_completion, handle_format_test, handle_generate, handle_render, handle_tokenize
+from handlers import handle_cache_prefill, handle_capabilities, handle_completion, handle_format_test, handle_generate, handle_render, handle_tokenize
 from handlers.cancel import request_cancel, reset_cancel
 
 
@@ -87,7 +87,25 @@ class Server:
                 )
 
             elif method == 'cache_prefill':
-                self._error_response("cache_prefill is not supported by the PyTorch backend")
+                cache_path = req.get('cache_path')
+                messages = req.get('messages')
+                if not cache_path or not messages:
+                    self._error_response("'cache_path' and 'messages' fields are required for cache_prefill")
+                    return
+                handle_cache_prefill(
+                    self.backend,
+                    self.capabilities,
+                    cache_path,
+                    messages,
+                    base_cache_path=req.get('base_cache_path'),
+                    trim_to_tokens=req.get('trim_to_tokens'),
+                    prefix_offsets=req.get('prefix_offsets'),
+                    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':
                 messages = req.get('messages')

package/src/pytorch/templates/cpu-minimal/tests/test_cache_handler.py

@@ -0,0 +1,284 @@
+import json
+from types import SimpleNamespace
+
+from handlers.cache import handle_cache_prefill
+from handlers.generate import handle_generate
+
+
+class _Tokenizer:
+    bos_token = "<bos>"
+    chat_template = "template"
+
+    def __init__(self):
+        self.calls = []
+
+    def apply_chat_template(self, messages, **kwargs):
+        self.calls.append((messages, kwargs))
+        return "prefix"
+
+
+class _Cache:
+    pass
+
+
+class _Backend:
+    model_kind = "lm"
+
+    def __init__(self, cache=None, multi_chunk=False, cache_write_tokens=0):
+        self.tokenizer = _Tokenizer()
+        self.cache = cache
+        self.cache_offsets = {cache: 2} if cache is not None else {}
+        self.multi_chunk = multi_chunk
+        self.pending_cache_write_tokens = cache_write_tokens
+        self.calls = []
+
+    def get_tokenizer(self):
+        return self.tokenizer
+
+    def cache_prefill(self, cache_path, prompt, **kwargs):
+        self.calls.append(("prefill", cache_path, prompt, kwargs))
+        self.pending_cache_write_tokens = 2
+        return {
+            "cache_path": cache_path,
+            "token_count": 2,
+            "cache_write_tokens": 2,
+        }
+
+    def consume_cache_write_tokens(self, cache_path):
+        self.calls.append(("consume-write", cache_path))
+        result = self.pending_cache_write_tokens
+        self.pending_cache_write_tokens = 0
+        return result
+
+    def load_cache_from_file(self, cache_path, **kwargs):
+        self.calls.append(("load", cache_path, kwargs))
+        return self.cache
+
+    def get_cache_offset(self, prompt_cache):
+        return self.cache_offsets.get(prompt_cache, 0)
+
+    def trim_cache(self, prompt_cache, tokens):
+        self.calls.append(("trim", prompt_cache, tokens))
+        trimmed_cache = _Cache()
+        self.cache_offsets[trimmed_cache] = max(
+            0,
+            self.cache_offsets[prompt_cache] - tokens,
+        )
+        return trimmed_cache
+
+    def tokenize_prompt(self, prompt):
+        assert prompt == "prefix suffix"
+        return [1, 2, 3]
+
+    def stream_generate(self, prompt, options, images=None, prompt_cache=None):
+        self.calls.append(("generate", prompt, options, images, prompt_cache))
+        cache_read_tokens = self.cache_offsets.get(prompt_cache)
+        if self.multi_chunk:
+            for index, text in enumerate(("a", "b", "c"), start=1):
+                yield SimpleNamespace(
+                    text=text,
+                    prompt_tokens=3 if index == 1 else None,
+                    generation_tokens=index,
+                    cache_read_tokens=cache_read_tokens if index == 1 else None,
+                )
+            return
+        yield SimpleNamespace(
+            text="ok",
+            prompt_tokens=3,
+            generation_tokens=1,
+            cache_read_tokens=cache_read_tokens,
+        )
+
+
+def _json_response(output):
+    return json.loads(output.split("\0", 1)[0])
+
+
+def test_cache_prefill_renders_without_generation_prompt(capsys):
+    backend = _Backend()
+
+    handle_cache_prefill(
+        backend,
+        {"special_tokens": {}},
+        "memory://prefix",
+        [{"role": "user", "content": "hello"}],
+    )
+
+    assert _json_response(capsys.readouterr().out) == {
+        "cache_path": "memory://prefix",
+        "token_count": 2,
+        "cache_write_tokens": 2,
+    }
+    assert backend.calls == [
+        ("prefill", "memory://prefix", "prefix", {
+            "base_cache_path": None,
+            "trim_to_tokens": None,
+            "prefix_offsets": None,
+            "prefix_hashes": None,
+            "images": None,
+            "max_image_size": 768,
+        })
+    ]
+    assert backend.tokenizer.calls[0][1]["add_generation_prompt"] is False
+
+
+def test_cache_prefill_passes_incremental_options_and_propagates_prefix_meta(capsys):
+    backend = _Backend()
+
+    handle_cache_prefill(
+        backend,
+        {"special_tokens": {}},
+        "tmp/extended.pytorch-cache",
+        [{"role": "user", "content": "hello"}],
+        base_cache_path="tmp/base.pytorch-cache",
+        trim_to_tokens=1,
+        prefix_offsets=[1, 2],
+        prefix_hashes=["hash-prefix", "hash-full"],
+    )
+
+    result = _json_response(capsys.readouterr().out)
+    assert result["prefix_offsets"] == [1, 2]
+    assert result["prefix_hashes"] == ["hash-prefix", "hash-full"]
+    assert backend.calls == [
+        ("prefill", "tmp/extended.pytorch-cache", "prefix", {
+            "base_cache_path": "tmp/base.pytorch-cache",
+            "trim_to_tokens": 1,
+            "prefix_offsets": [1, 2],
+            "prefix_hashes": ["hash-prefix", "hash-full"],
+            "images": None,
+            "max_image_size": 768,
+        })
+    ]
+
+
+def test_generate_loads_cache_and_only_generates_suffix(capsys):
+    cache = _Cache()
+    backend = _Backend(cache, cache_write_tokens=2)
+
+    handle_generate(
+        backend,
+        "prefix suffix",
+        options={"max_tokens": 1},
+        cache_path="memory://prefix",
+    )
+
+    output = capsys.readouterr().out
+    assert output.startswith("ok")
+    meta = json.loads(output.split("\x1e__META__:", 1)[1].split("\0", 1)[0])
+    assert meta == {
+        "prompt_tokens": 3,
+        "generation_tokens": 1,
+        "cache_read_tokens": 2,
+        "cache_write_tokens": 2,
+        "cache_loaded": True,
+    }
+    generate_call = next(call for call in backend.calls if call[0] == "generate")
+    assert generate_call[1] == [3]
+    assert generate_call[4] is cache
+    load_call = next(call for call in backend.calls if call[0] == "load")
+    assert load_call[2]["prompt"] == "prefix suffix"
+    assert load_call[2]["prefix_token_count"] is None
+    assert load_call[2]["images"] is None
+    assert load_call[2]["max_image_size"] == 768
+    assert ("consume-write", "memory://prefix") in backend.calls
+
+
+def test_generate_trims_loaded_cache_before_generating(capsys):
+    cache = _Cache()
+    backend = _Backend(cache, cache_write_tokens=2)
+
+    handle_generate(
+        backend,
+        "prefix suffix",
+        options={"max_tokens": 1},
+        cache_path="memory://prefix",
+        cache_trim_tokens=1,
+    )
+
+    output = capsys.readouterr().out
+    generate_call = next(call for call in backend.calls if call[0] == "generate")
+    assert generate_call[1] == [2, 3]
+    assert generate_call[4] is not cache
+    assert backend.get_cache_offset(cache) == 2
+    assert backend.get_cache_offset(generate_call[4]) == 1
+    assert ("trim", cache, 1) in backend.calls
+    load_call = next(call for call in backend.calls if call[0] == "load")
+    assert load_call[2]["prefix_token_count"] == 1
+    meta = json.loads(output.split("\x1e__META__:", 1)[1].split("\0", 1)[0])
+    assert meta["cache_read_tokens"] == 1
+
+    # The original ref remains usable after the trimmed generation.
+    handle_generate(
+        backend,
+        "prefix suffix",
+        options={"max_tokens": 1},
+        cache_path="memory://prefix",
+    )
+    second_output = capsys.readouterr().out
+    second_generate_call = [
+        call for call in backend.calls if call[0] == "generate"
+    ][-1]
+    assert second_generate_call[1] == [3]
+    assert second_generate_call[4] is cache
+    second_meta = json.loads(
+        second_output.split("\x1e__META__:", 1)[1].split("\0", 1)[0]
+    )
+    assert second_meta["cache_read_tokens"] == 2
+
+
+def test_generate_preserves_usage_meta_across_multiple_chunks(capsys):
+    cache = _Cache()
+    backend = _Backend(cache, multi_chunk=True, cache_write_tokens=2)
+
+    handle_generate(
+        backend,
+        "prefix suffix",
+        options={"max_tokens": 3},
+        cache_path="memory://prefix",
+    )
+
+    output = capsys.readouterr().out
+    assert output.startswith("abc")
+    meta = json.loads(output.split("\x1e__META__:", 1)[1].split("\0", 1)[0])
+    assert meta == {
+        "prompt_tokens": 3,
+        "generation_tokens": 3,
+        "cache_read_tokens": 2,
+        "cache_write_tokens": 2,
+        "cache_loaded": True,
+    }
+
+
+def test_generate_uses_cold_path_for_missing_cache(capsys):
+    backend = _Backend(cache=None)
+
+    handle_generate(
+        backend,
+        "prefix suffix",
+        cache_path="memory://missing",
+    )
+
+    output = capsys.readouterr().out
+    meta = json.loads(output.split("\x1e__META__:", 1)[1].split("\0", 1)[0])
+    assert meta["cache_loaded"] is False
+    assert "cache_read_tokens" not in meta
+    assert "cache_write_tokens" not in meta
+    generate_call = next(call for call in backend.calls if call[0] == "generate")
+    assert generate_call[1] == "prefix suffix"
+    assert generate_call[4] is None
+
+
+def test_generate_removes_cached_prefix_from_token_ids(capsys):
+    cache = _Cache()
+    backend = _Backend(cache)
+
+    handle_generate(
+        backend,
+        [1, 2, 3],
+        cache_path="memory://prefix",
+    )
+
+    output = capsys.readouterr().out
+    assert '"cache_loaded": true' in output
+    generate_call = next(call for call in backend.calls if call[0] == "generate")
+    assert generate_call[1] == [3]

package/src/pytorch/templates/cpu-minimal/tests/test_capabilities.py

@@ -0,0 +1,14 @@
+from utils.token_utils import get_capabilities
+
+
+class _Tokenizer:
+    apply_chat_template = None
+    chat_template = None
+    special_tokens_map = {}
+    added_tokens_encoder = {}
+
+
+def test_capabilities_advertise_cache_prefill():
+    capabilities = get_capabilities(_Tokenizer())
+
+    assert "cache_prefill" in capabilities["methods"]