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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: cortexgrid-infer
-Version: 0.2.4
+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
@@ -184,16 +184,25 @@ Both serve apps are tuned rather than defaulted:
   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
-it, and a host outrunning its accelerator is a GPU problem to begin with.
+  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.4 → cortexgrid_infer-0.2.6}/README.md

@@ -158,16 +158,25 @@ Both serve apps are tuned rather than defaulted:
   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
-it, and a host outrunning its accelerator is a GPU problem to begin with.
+  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.4 → cortexgrid_infer-0.2.6}/cortexgrid_infer/providers/huggingface_complete_serve.py

@@ -29,8 +29,14 @@ 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
+# 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",
@@ -55,8 +61,12 @@ class HuggingFaceCompletingDeployment:
             str(path), torch_dtype=torch.float16
         )
         self._model.to(self._device)  # type: ignore
-        self._max_input_tokens = getattr(
-            self._model.config, "max_position_embeddings", FALLBACK_MAX_INPUT_TOKENS
+        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
@@ -66,27 +76,39 @@ class HuggingFaceCompletingDeployment:
         if self._compiled:
             config = self._model.generation_config
             config.cache_implementation = "static"
-            config.max_cache_len = self._max_input_tokens
+            config.max_cache_len = self._max_total_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."""
+    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 <= self._max_input_tokens:
+        if length <= limit:
             return inputs
         log.warning(
-            "truncating prompt from %d to %d tokens, the model's context",
-            length, self._max_input_tokens,
+            "truncating prompt from %d to %d tokens, the room this replica has",
+            length, limit,
         )
         for key, value in inputs.items():
-            inputs[key] = value[..., -self._max_input_tokens :]
+            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)
@@ -99,9 +121,10 @@ class HuggingFaceCompletingDeployment:
             tokenize=False,
             add_generation_prompt=True,
         )
-        inputs = self._truncate(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.4 → cortexgrid_infer-0.2.6}/pyproject.toml

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