cortexgrid-infer 0.2.1__tar.gz → 0.2.2__tar.gz

{cortexgrid_infer-0.2.1 → cortexgrid_infer-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: cortexgrid-infer
-Version: 0.2.1
+Version: 0.2.2
 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
@@ -19,7 +19,7 @@ Requires-Dist: protobuf>=4.0
 Requires-Dist: python-dotenv>=1.0.0
 Requires-Dist: sentencepiece>=0.2
 Requires-Dist: torch>=2.0.0
-Requires-Dist: transformers>=4.40.0
+Requires-Dist: transformers>=4.47.0
 Provides-Extra: dev
 Requires-Dist: pytest>=8.0; extra == 'dev'
 Description-Content-Type: text/markdown
@@ -149,6 +149,58 @@ the whole model onto a single device. A model too large for any one card needs a
 app that shards it, not a larger `num_gpus`: the requirement would simply never be
 placed. (`AnthropicImport` asks for none — it holds no weights.)
 
+## Compilation
+
+Both HuggingFace serve apps run `torch.compile` over the model they load, because
+eager PyTorch bills per *operation* and a transformer forward is mostly glue —
+views, broadcasts, permutes, elementwise adds — with only a small fraction of the
+dispatches being the matrix multiplies that do the arithmetic. While the tensors
+are large the host stays ahead of the accelerator and none of that shows; once
+they are small, the accelerator drains its queue faster than the host can fill it
+and step time tracks the operation count rather than the work. Feeding the model
+less then stops helping, because the op count does not shrink with the input.
+Inductor fuses the glue away; `reduce-overhead` additionally captures a CUDA graph,
+so a step replays one graph instead of thousands of launches.
+
+`cortexgrid_infer.compiling` offers two modes, and a serve app picks the one that
+fits the work it does:
+
+| | | for |
+|---|---|---|
+| `Mode.GRAPHED` (`reduce-overhead`) | fuses **and** captures a CUDA graph, so a forward replays one graph instead of issuing its launches | a loop of identically shaped forwards, where the capture is reused every iteration. A graph cannot go dynamic, so each distinct shape captures separately |
+| `Mode.FUSED` (`default`) | fuses only | work whose shapes move under it, where there is nothing stable to capture and dynamo is free to compile one dynamic kernel set for all of them |
+
+Both serve apps are tuned rather than defaulted:
+
+- **`HuggingFaceImageDeployment`** compiles the denoiser (`transformer` or `unet`)
+  and every `text_encoder*` as `Mode.GRAPHED`. A denoising schedule is the ideal case —
+  every step is the same shapes — so the capture is made once and replayed for the
+  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.
+
+Two things worth knowing:
+
+- **Warm-up is per input shape, and lands on a request.** Compilation is lazy, so
+  the first request at each image size — or each prompt and generation length —
+  pays for it.
+- **Compilation never becomes load-bearing.** A backend that cannot compile a model
+  falls back to eager rather than failing the deployment.
+
+It does not fix everything. Every forward re-reads the model's whole weight tensor
+from memory, and no amount of fusion makes that read smaller. Below some input size
+it is the floor; only quantizing the weights or running fewer forwards moves it.
+
 ## Inference
 
 `complete` streams `CompletionChunk`s (`.content`, `.tool_calls`, `.finish_reason`):
@@ -215,6 +267,8 @@ serve app re-encodes them into the text form the client parses.
 | `HuggingFaceImageModel(url, model_id)` | Client for a deployed diffusers app. |
 | `split_model_id(model_id) -> (family, suffix)` | The registry identity an importer derives. |
 | `detect_device()` | The torch device a serve app should load onto. |
+| `compiling.Mode.GRAPHED` / `compiling.Mode.FUSED` | The two modes a serve app chooses between. See [Compilation](#compilation). |
+| `compiling.supported(device)` / `compile_module(module, device, mode)` / `compile_pipeline(pipe, device, mode)` | Whether to compile here, and what a serve app calls to do it. |
 | `ModelDeployFailed` | Re-exported from cortexgrid; subclasses `RuntimeError`. |
 
 Types: `DeployedModel`, `CompletingModel`, `GeneratingModel`, `CompletionChunk`,

{cortexgrid_infer-0.2.1 → cortexgrid_infer-0.2.2}/README.md

@@ -123,6 +123,58 @@ the whole model onto a single device. A model too large for any one card needs a
 app that shards it, not a larger `num_gpus`: the requirement would simply never be
 placed. (`AnthropicImport` asks for none — it holds no weights.)
 
+## Compilation
+
+Both HuggingFace serve apps run `torch.compile` over the model they load, because
+eager PyTorch bills per *operation* and a transformer forward is mostly glue —
+views, broadcasts, permutes, elementwise adds — with only a small fraction of the
+dispatches being the matrix multiplies that do the arithmetic. While the tensors
+are large the host stays ahead of the accelerator and none of that shows; once
+they are small, the accelerator drains its queue faster than the host can fill it
+and step time tracks the operation count rather than the work. Feeding the model
+less then stops helping, because the op count does not shrink with the input.
+Inductor fuses the glue away; `reduce-overhead` additionally captures a CUDA graph,
+so a step replays one graph instead of thousands of launches.
+
+`cortexgrid_infer.compiling` offers two modes, and a serve app picks the one that
+fits the work it does:
+
+| | | for |
+|---|---|---|
+| `Mode.GRAPHED` (`reduce-overhead`) | fuses **and** captures a CUDA graph, so a forward replays one graph instead of issuing its launches | a loop of identically shaped forwards, where the capture is reused every iteration. A graph cannot go dynamic, so each distinct shape captures separately |
+| `Mode.FUSED` (`default`) | fuses only | work whose shapes move under it, where there is nothing stable to capture and dynamo is free to compile one dynamic kernel set for all of them |
+
+Both serve apps are tuned rather than defaulted:
+
+- **`HuggingFaceImageDeployment`** compiles the denoiser (`transformer` or `unet`)
+  and every `text_encoder*` as `Mode.GRAPHED`. A denoising schedule is the ideal case —
+  every step is the same shapes — so the capture is made once and replayed for the
+  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.
+
+Two things worth knowing:
+
+- **Warm-up is per input shape, and lands on a request.** Compilation is lazy, so
+  the first request at each image size — or each prompt and generation length —
+  pays for it.
+- **Compilation never becomes load-bearing.** A backend that cannot compile a model
+  falls back to eager rather than failing the deployment.
+
+It does not fix everything. Every forward re-reads the model's whole weight tensor
+from memory, and no amount of fusion makes that read smaller. Below some input size
+it is the floor; only quantizing the weights or running fewer forwards moves it.
+
 ## Inference
 
 `complete` streams `CompletionChunk`s (`.content`, `.tool_calls`, `.finish_reason`):
@@ -189,6 +241,8 @@ serve app re-encodes them into the text form the client parses.
 | `HuggingFaceImageModel(url, model_id)` | Client for a deployed diffusers app. |
 | `split_model_id(model_id) -> (family, suffix)` | The registry identity an importer derives. |
 | `detect_device()` | The torch device a serve app should load onto. |
+| `compiling.Mode.GRAPHED` / `compiling.Mode.FUSED` | The two modes a serve app chooses between. See [Compilation](#compilation). |
+| `compiling.supported(device)` / `compile_module(module, device, mode)` / `compile_pipeline(pipe, device, mode)` | Whether to compile here, and what a serve app calls to do it. |
 | `ModelDeployFailed` | Re-exported from cortexgrid; subclasses `RuntimeError`. |
 
 Types: `DeployedModel`, `CompletingModel`, `GeneratingModel`, `CompletionChunk`,

cortexgrid_infer-0.2.2/cortexgrid_infer/compiling.py

@@ -0,0 +1,70 @@
+"""`torch.compile` for the serve apps that run a model in-process."""
+
+from __future__ import annotations
+
+import logging
+from enum import StrEnum
+from typing import Any
+
+log = logging.getLogger(__name__)
+
+# Inductor is weakest off CUDA, and a host outrunning its accelerator is a GPU problem.
+_ACCELERATOR = "cuda"
+
+# Parts of a diffusers pipeline worth compiling. Not the VAE: smallest win, shakiest shapes.
+_DENOISER_ATTRS = ("transformer", "unet")
+_TEXT_ENCODER_ATTRS = ("text_encoder", "text_encoder_2", "text_encoder_3")
+
+
+class Mode(StrEnum):
+    """How far `torch.compile` should go. A serve app picks by the shape of its work."""
+
+    GRAPHED = "reduce-overhead"
+    """Fuse, then capture a CUDA graph. For a loop of identically shaped forwards;
+    captures per shape, so it wants a settled set of them."""
+
+    FUSED = "default"
+    """Fuse only. For work whose shapes move under it, leaving dynamo free to
+    compile one dynamic kernel set for all of them."""
+
+
+def supported(device: Any) -> bool:
+    """Whether compiling is worth doing on this device."""
+    return getattr(device, "type", None) == _ACCELERATOR
+
+
+def compile_module(module: Any, device: Any, mode: Mode, label: str = "model") -> bool:
+    """Compile one module's forward, in place. Returns whether it was compiled.
+
+    Lazily, as `torch.compile` always is: the first forward through each input
+    shape is the one that pays."""
+    if module is None or not supported(device):
+        return False
+    try:
+        _fall_back_to_eager_on_failure()
+        module.compile(mode=mode)
+    except Exception as exc:  # noqa: BLE001 - an optimisation, never fatal
+        log.warning("torch.compile(%s, mode=%s) failed: %s", label, mode, exc)
+        return False
+    log.info("compiling %s with torch.compile(mode=%s)", label, mode)
+    return True
+
+
+def compile_pipeline(pipe: Any, device: Any, mode: Mode) -> list[str]:
+    """Compile a diffusers pipeline's denoiser and text encoders. Returns which."""
+    compiled = []
+    for attr in (*_DENOISER_ATTRS, *_TEXT_ENCODER_ATTRS):
+        module = getattr(pipe, attr, None)
+        if module is not None and compile_module(module, device, mode, label=attr):
+            compiled.append(attr)
+    return compiled
+
+
+def _fall_back_to_eager_on_failure() -> None:
+    """Make dynamo warn rather than raise out of a request. Process-wide."""
+    try:
+        import torch._dynamo.config as dynamo_config
+
+        dynamo_config.suppress_errors = True
+    except Exception:  # noqa: BLE001 - best effort on a private flag
+        pass

{cortexgrid_infer-0.2.1 → cortexgrid_infer-0.2.2}/cortexgrid_infer/providers/huggingface_complete_serve.py

@@ -11,9 +11,15 @@ 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
 
 
@@ -43,6 +49,14 @@ class HuggingFaceCompletingDeployment:
             str(path), torch_dtype=torch.float16
         )
         self._model.to(self._device)  # type: ignore
+        # Decode has no stable shape to capture until the cache is static; given one,
+        # transformers compiles `generate` itself, so compiling the module would double up.
+        self._compiled = compiling.supported(self._device)
+        if self._compiled:
+            self._model.generation_config.cache_implementation = "static"
+            self._model.generation_config.compile_config = CompileConfig(
+                mode=compiling.Mode.GRAPHED
+            )
 
     @_app.post("/complete")
     async def complete(self, body: dict[str, Any]) -> StreamingResponse:

{cortexgrid_infer-0.2.1 → cortexgrid_infer-0.2.2}/cortexgrid_infer/providers/huggingface_image_serve.py

@@ -25,6 +25,7 @@ from fastapi import FastAPI
 from PIL import Image, ImageOps
 import torch
 
+from cortexgrid_infer import compiling
 from cortexgrid_infer.device import detect_device
 
 
@@ -70,6 +71,10 @@ class HuggingFaceImageDeployment:
         if vae is not None and hasattr(vae, "enable_tiling"):
             vae.enable_tiling()
         self._pipe = pipe
+        # A denoising schedule is a loop of identically shaped forwards: capture once, replay.
+        self._compiled = compiling.compile_pipeline(
+            pipe, self._device, compiling.Mode.GRAPHED
+        )
 
     @_app.post("/generate")
     async def generate(self, body: dict[str, Any]) -> dict[str, Any]:

{cortexgrid_infer-0.2.1 → cortexgrid_infer-0.2.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "cortexgrid-infer"
-version = "0.2.1"
+version = "0.2.2"
 description = "Deploy and interact with models across providers (Anthropic, HuggingFace LLM + diffusion)"
 readme = "README.md"
 license = "Apache-2.0"
@@ -17,7 +17,9 @@ dependencies = [
     # redundant PUT - which is why this library has no reuse guard of its own.
     "cortexgrid>=0.3.4",
     "torch>=2.0.0",
-    "transformers>=4.40.0",
+    # 4.47: CompileConfig, which the completion serve app hands to `generate`
+    # alongside a static cache so transformers compiles the decode step itself.
+    "transformers>=4.47.0",
     "huggingface_hub>=0.20.0",
     "anthropic>=0.40.0",
     "python-dotenv>=1.0.0",