cortexgrid-infer 0.2.3__tar.gz → 0.2.4__tar.gz

{cortexgrid_infer-0.2.3 → cortexgrid_infer-0.2.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: cortexgrid-infer
-Version: 0.2.3
+Version: 0.2.4
 Summary: Deploy and interact with models across providers (Anthropic, HuggingFace LLM + diffusion)
 Project-URL: Homepage, https://github.com/robodatalab/cortexgrid-infer
 Project-URL: Repository, https://github.com/robodatalab/cortexgrid-infer
@@ -178,12 +178,18 @@ Both serve apps are tuned rather than defaulted:
   whole loop. The VAE is left alone: smallest win of the three, and the one whose
   shapes move most, especially with tiling on.
 - **`HuggingFaceCompletingDeployment`** does *not* compile the module. Decode has
-  no stable shape until the KV cache is static, so it sets
-  `cache_implementation="static"` and a `CompileConfig(mode=Mode.GRAPHED)`, and
-  transformers compiles `generate` itself — capturing prefill and decode
-  separately, which it is better placed to do. Compiling the module here as well
-  would only compile the same forward twice. The cost is VRAM: a static cache is
-  preallocated to the length asked for, where a dynamic one grows into it.
+  no stable shape of its own — the sequence grows a token per forward — so it
+  pins one: `cache_implementation="static"` with a fixed `max_cache_len`, plus a
+  `CompileConfig(mode=Mode.GRAPHED)`. transformers then compiles `generate`
+  itself, capturing prefill and decode separately, which it is better placed to
+  do; compiling the module here as well would compile the same forward twice.
+
+  The fixed length is the point, and it comes from the model: the cache is
+  pinned to `config.max_position_embeddings`. A prompt longer than that is
+  truncated to it — keeping the end, since a chat template puts the turn to
+  answer last — with a warning naming both lengths. Serving it instead would
+  resize the cache and recompile the decode loop, which costs more than the
+  generation itself.
 
 Neither asks to be configured, for the same reason the serve apps do not ask which
 dtype to load in. Both happen on CUDA and nowhere else — inductor is weakest off

{cortexgrid_infer-0.2.3 → cortexgrid_infer-0.2.4}/README.md

@@ -152,12 +152,18 @@ Both serve apps are tuned rather than defaulted:
   whole loop. The VAE is left alone: smallest win of the three, and the one whose
   shapes move most, especially with tiling on.
 - **`HuggingFaceCompletingDeployment`** does *not* compile the module. Decode has
-  no stable shape until the KV cache is static, so it sets
-  `cache_implementation="static"` and a `CompileConfig(mode=Mode.GRAPHED)`, and
-  transformers compiles `generate` itself — capturing prefill and decode
-  separately, which it is better placed to do. Compiling the module here as well
-  would only compile the same forward twice. The cost is VRAM: a static cache is
-  preallocated to the length asked for, where a dynamic one grows into it.
+  no stable shape of its own — the sequence grows a token per forward — so it
+  pins one: `cache_implementation="static"` with a fixed `max_cache_len`, plus a
+  `CompileConfig(mode=Mode.GRAPHED)`. transformers then compiles `generate`
+  itself, capturing prefill and decode separately, which it is better placed to
+  do; compiling the module here as well would compile the same forward twice.
+
+  The fixed length is the point, and it comes from the model: the cache is
+  pinned to `config.max_position_embeddings`. A prompt longer than that is
+  truncated to it — keeping the end, since a chat template puts the turn to
+  answer last — with a warning naming both lengths. Serving it instead would
+  resize the cache and recompile the decode loop, which costs more than the
+  generation itself.
 
 Neither asks to be configured, for the same reason the serve apps do not ask which
 dtype to load in. Both happen on CUDA and nowhere else — inductor is weakest off

{cortexgrid_infer-0.2.3 → cortexgrid_infer-0.2.4}/cortexgrid_infer/providers/huggingface_complete_serve.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import asyncio
+import logging
 from collections.abc import AsyncIterator
 from threading import Thread
 from typing import Any
@@ -11,16 +12,26 @@ import cortexgrid
 from cortexgrid import serve
 from fastapi import FastAPI
 from fastapi.responses import StreamingResponse
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    CompileConfig,
+    TextIteratorStreamer,
+)
 import torch
 
 from cortexgrid_infer import compiling
 from cortexgrid_infer.device import detect_device
 
 
+log = logging.getLogger(__name__)
+
 _app = FastAPI()
 
 
+# Used only by a model whose config states no context length of its own.
+FALLBACK_MAX_INPUT_TOKENS = 4096
+
 _RESERVED_BODY_KEYS = {
     "messages",
     "tools",
@@ -44,10 +55,32 @@ class HuggingFaceCompletingDeployment:
             str(path), torch_dtype=torch.float16
         )
         self._model.to(self._device)  # type: ignore
-        # Decode's shapes move as the sequence grows, so fuse without capturing.
-        self._compiled = compiling.compile_module(
-            self._model, self._device, compiling.Mode.FUSED, label="model"
+        self._max_input_tokens = getattr(
+            self._model.config, "max_position_embeddings", FALLBACK_MAX_INPUT_TOKENS
         )
+        # A static cache of fixed length is what makes decode compilable: it
+        # pins the shape every forward sees, so the loop is captured once
+        # instead of recompiled per token. transformers then compiles `generate`
+        # itself, which is why the module is not compiled here as well.
+        self._compiled = compiling.supported(self._device)
+        if self._compiled:
+            config = self._model.generation_config
+            config.cache_implementation = "static"
+            config.max_cache_len = self._max_input_tokens
+            config.compile_config = CompileConfig(mode=compiling.Mode.GRAPHED)
+
+    def _truncate(self, inputs: Any) -> Any:
+        """Cut an over-long prompt to what the model handles, keeping its end."""
+        length = inputs["input_ids"].shape[-1]
+        if length <= self._max_input_tokens:
+            return inputs
+        log.warning(
+            "truncating prompt from %d to %d tokens, the model's context",
+            length, self._max_input_tokens,
+        )
+        for key, value in inputs.items():
+            inputs[key] = value[..., -self._max_input_tokens :]
+        return inputs
 
     @_app.post("/complete")
     async def complete(self, body: dict[str, Any]) -> StreamingResponse:
@@ -66,7 +99,9 @@ class HuggingFaceCompletingDeployment:
             tokenize=False,
             add_generation_prompt=True,
         )
-        inputs = self._tokenizer(prompt, return_tensors="pt").to(self._model.device)
+        inputs = self._truncate(self._tokenizer(prompt, return_tensors="pt")).to(
+            self._model.device
+        )
 
         queue: asyncio.Queue[str | None] = asyncio.Queue()
         error: list[BaseException] = []

{cortexgrid_infer-0.2.3 → cortexgrid_infer-0.2.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "cortexgrid-infer"
-version = "0.2.3"
+version = "0.2.4"
 description = "Deploy and interact with models across providers (Anthropic, HuggingFace LLM + diffusion)"
 readme = "README.md"
 license = "Apache-2.0"