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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: cortexgrid-infer
-Version: 0.2.3
+Version: 0.2.6
 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,16 +178,31 @@ 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.
-
-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
-it, and a host outrunning its accelerator is a GPU problem to begin with.
+  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 how long it is decides what decode costs:
+  every token's attention reads the whole cache, filled or not. That is why it is
+  *not* pinned to the model's context — 32k slots of a Qwen2.5-3B KV cache is
+  1.2 GB of keys and values re-read per generated token, several times the
+  traffic of the model's own weights, to make room for a prompt nobody sent. It
+  is pinned instead to the prompt and reply a caller plausibly sends: 4096 tokens,
+  or `max_total_tokens` on the model card, never past the model's own
+  `max_position_embeddings`. A reply takes at most half of it, so an
+  over-optimistic `max_new_tokens` cannot squeeze the prompt to nothing, and a
+  prompt that will not fit in the rest is truncated — keeping the end, since a
+  chat template puts the turn to answer last — with a warning naming both
+  lengths. Serving it whole would resize the cache and recompile the decode loop,
+  which costs more than the generation itself.
+
+The image app asks to be configured for nothing, for the same reason neither asks
+which dtype to load in; the completion app takes `max_total_tokens` alone, because
+how much cache a caller needs is the one thing the model cannot tell it. Both
+happen on CUDA and nowhere else — inductor is weakest off it, and a host
+outrunning its accelerator is a GPU problem to begin with.
 
 Two things worth knowing:
 

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

@@ -152,16 +152,31 @@ 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.
-
-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
-it, and a host outrunning its accelerator is a GPU problem to begin with.
+  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 how long it is decides what decode costs:
+  every token's attention reads the whole cache, filled or not. That is why it is
+  *not* pinned to the model's context — 32k slots of a Qwen2.5-3B KV cache is
+  1.2 GB of keys and values re-read per generated token, several times the
+  traffic of the model's own weights, to make room for a prompt nobody sent. It
+  is pinned instead to the prompt and reply a caller plausibly sends: 4096 tokens,
+  or `max_total_tokens` on the model card, never past the model's own
+  `max_position_embeddings`. A reply takes at most half of it, so an
+  over-optimistic `max_new_tokens` cannot squeeze the prompt to nothing, and a
+  prompt that will not fit in the rest is truncated — keeping the end, since a
+  chat template puts the turn to answer last — with a warning naming both
+  lengths. Serving it whole would resize the cache and recompile the decode loop,
+  which costs more than the generation itself.
+
+The image app asks to be configured for nothing, for the same reason neither asks
+which dtype to load in; the completion app takes `max_total_tokens` alone, because
+how much cache a caller needs is the one thing the model cannot tell it. Both
+happen on CUDA and nowhere else — inductor is weakest off it, and a host
+outrunning its accelerator is a GPU problem to begin with.
 
 Two things worth knowing:
 

{cortexgrid_infer-0.2.3 → cortexgrid_infer-0.2.6}/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,32 @@ 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()
 
 
+# Prompt + reply a replica is sized for. The static cache is allocated at this
+# many slots and every decoded token reads all of them, filled or not: at
+# Qwen2.5-3B's 32k context that is 1.2 GB of keys and values re-read per token,
+# which costs far more than the prompt it makes room for. So the default is what
+# a caller plausibly sends rather than what the model could hold; a model that
+# needs more takes `max_total_tokens` on its model card, up to its own context.
+MAX_TOTAL_TOKENS_PARAM = "max_total_tokens"
+DEFAULT_MAX_TOTAL_TOKENS = 4096
+
 _RESERVED_BODY_KEYS = {
     "messages",
     "tools",
@@ -44,16 +61,54 @@ 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"
+        settings = cortexgrid.model_config(family, suffix, run_name)
+        context = getattr(
+            self._model.config, "max_position_embeddings", DEFAULT_MAX_TOTAL_TOKENS
+        )
+        self._max_total_tokens = min(
+            int(settings.get(MAX_TOTAL_TOKENS_PARAM, DEFAULT_MAX_TOTAL_TOKENS)), context
         )
+        # 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_total_tokens
+            config.compile_config = CompileConfig(mode=compiling.Mode.GRAPHED)
+
+    def _room_for_the_reply(self, requested: int) -> int:
+        """How many tokens a reply may take, leaving the prompt the rest.
+
+        Half the budget at most, so a caller asking for more output than the
+        replica is sized for cannot squeeze the prompt down to nothing."""
+        return min(requested, self._max_total_tokens // 2)
+
+    def _truncate(self, inputs: Any, limit: int) -> Any:
+        """Cut an over-long prompt to ``limit`` tokens, keeping its end.
+
+        A prompt and reply that together outgrew the cache would reallocate it
+        and recompile the decode loop, which costs more than the generation."""
+        length = inputs["input_ids"].shape[-1]
+        if length <= limit:
+            return inputs
+        log.warning(
+            "truncating prompt from %d to %d tokens, the room this replica has",
+            length, limit,
+        )
+        for key, value in inputs.items():
+            inputs[key] = value[..., -limit:]
+        return inputs
 
     @_app.post("/complete")
     async def complete(self, body: dict[str, Any]) -> StreamingResponse:
         messages = body["messages"]
         tools = body.get("tools")
-        max_new_tokens = body.get("max_new_tokens", 16 * 1024)
+        max_new_tokens = self._room_for_the_reply(
+            body.get("max_new_tokens", self._max_total_tokens)
+        )
         temperature = body.get("temperature", 0.7)
         top_p = body.get("top_p", 0.9)
         top_k = body.get("top_k", 50)
@@ -66,7 +121,10 @@ 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"),
+            self._max_total_tokens - max_new_tokens,
+        ).to(self._model.device)
 
         queue: asyncio.Queue[str | None] = asyncio.Queue()
         error: list[BaseException] = []

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

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