From f64a12237797d49d5bfea239a1c2d89e48c02cec Mon Sep 17 00:00:00 2001
From: Amrit Krishnan <amrit110@gmail.com>
Date: Sat, 17 Oct 2026 15:32:38 -0400
Subject: [PATCH 01/39] feat(tabicl): compressed ICL KV cache (fp16/int8) +
 memory benchmarks

backends/tabicl_kv.py: CompressedKVEntry duck-types tabicl KVCacheEntry,
storing the ICL cache int8 (per head/row scale, ~0.52x of fp16) or fp16
(small tables kept fp32 by tabicl), decompressing one layer at a time on
access. TABCTX_KV_PRECISION=native|fp16|int8 (default native until the
accuracy study). Benchmarks: bench_memory_pressure (min working memory
per shape), eval_kv_precision (accuracy/memory/latency), profiler now
records TabICL cache-mode stages.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 benchmarks/bench_memory_pressure.py | 133 ++++++++++++++++++++++++
 benchmarks/eval_kv_precision.py     | 141 ++++++++++++++++++++++++++
 benchmarks/profile_fit_predict.py   |  69 ++++++++-----
 src/tabctx/backends/kv_precision.py |   9 ++
 src/tabctx/backends/tabicl.py       |  15 ++-
 src/tabctx/backends/tabicl_kv.py    | 152 ++++++++++++++++++++++++++++
 src/tabctx/serve/factory.py         |  21 +++-
 tests/unit/test_serve_factory.py    |  12 +++
 tests/unit/test_tabicl_kv.py        | 107 ++++++++++++++++++++
 9 files changed, 629 insertions(+), 30 deletions(-)
 create mode 100644 benchmarks/bench_memory_pressure.py
 create mode 100644 benchmarks/eval_kv_precision.py
 create mode 100644 src/tabctx/backends/kv_precision.py
 create mode 100644 src/tabctx/backends/tabicl_kv.py
 create mode 100644 tests/unit/test_tabicl_kv.py

diff --git a/benchmarks/bench_memory_pressure.py b/benchmarks/bench_memory_pressure.py
new file mode 100644
index 0000000..6294926
--- /dev/null
+++ b/benchmarks/bench_memory_pressure.py
@@ -0,0 +1,133 @@
+#!/usr/bin/env python3
+"""How much GPU memory does a fit actually NEED? (#31, #37)
+
+TabICL's InferenceManager sizes each stage's batches from the GPU memory
+that is currently free (x a safety factor), so an unconstrained fit's peak
+expands to fill the card: it measures "what was free", not "what is
+needed". tabctx's admission learns from measured peaks, so it can reject
+tables that would fit.
+
+This runs the same fit with only F GiB left free (the rest taken by a
+ballast tensor) for decreasing F, and records success, fit time, transient
+peak and resident context size. The smallest F that still succeeds at an
+acceptable time is the real requirement.
+
+Usage (on a GPU host):
+  bench_memory_pressure.py --backend tabicl --shapes 20000x100 100000x10 \
+      --free-gib 0 24 16 8 4 2 --out pressure-tabicl.jsonl
+(free 0 = no ballast, the whole card)
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+import time
+
+import torch
+
+sys.path.insert(0, __file__.rsplit("/", 1)[0])
+from profile_fit_predict import build_backend, make_table, parse_shape  # noqa: E402
+
+GiB = 1024**3
+
+
+def run(args, backend, rows: int, features: int, free_gib: float) -> dict:
+    device = torch.device(args.device)
+    record = {"rows": rows, "features": features, "free_gib": free_gib or "all"}
+    ballast = None
+    try:
+        torch.cuda.empty_cache()
+        if free_gib:
+            free_now, _ = torch.cuda.mem_get_info(device)
+            take = int(free_now - free_gib * GiB)
+            if take <= 0:
+                record["status"] = "skip (less free than requested)"
+                return record
+            ballast = torch.empty(take, dtype=torch.uint8, device=device)
+        X, y = make_table(rows, features, args.task, seed=rows + features)
+        torch.cuda.synchronize(device)
+        t0 = time.perf_counter()
+        payload = backend.fit(X, y, args.task)
+        torch.cuda.synchronize(device)
+        record.update(
+            status="ok",
+            fit_s=round(time.perf_counter() - t0, 3),
+            peak_gib=round((backend.fit_peak_bytes_hint() or 0) / GiB, 2),
+            resident_gib=round(
+                (backend.context_bytes_hint(rows, features) or 0) / GiB, 2
+            ),
+        )
+        del payload
+    except Exception as e:  # noqa: BLE001 -- OOM is the result being measured
+        record.update(status=type(e).__name__, error=str(e)[:200])
+    finally:
+        del ballast
+        torch.cuda.empty_cache()
+    return record
+
+
+def search(args, backend, rows: int, features: int) -> dict:
+    """Binary-search the least free memory at which the fit still succeeds
+    (to --precision-gib), recording the fit time there and unconstrained."""
+    free_all = run(args, backend, rows, features, 0)
+    if free_all["status"] != "ok":
+        return {"rows": rows, "features": features, "status": "fails unconstrained"}
+    lo = free_all["resident_gib"]  # can't succeed with less than the context
+    hi = torch.cuda.mem_get_info(torch.device(args.device))[0] / GiB
+    best = None
+    while hi - lo > args.precision_gib:
+        mid = (lo + hi) / 2
+        r = run(args, backend, rows, features, round(mid, 2))
+        if r["status"] == "ok":
+            hi, best = mid, r
+        else:
+            lo = mid
+    return {
+        "rows": rows,
+        "features": features,
+        "status": "ok",
+        "resident_gib": free_all["resident_gib"],
+        "unconstrained_peak_gib": free_all["peak_gib"],
+        "unconstrained_fit_s": free_all["fit_s"],
+        "min_free_gib": round(hi, 2),
+        "min_working_gib": round(hi - free_all["resident_gib"], 2),
+        "fit_s_at_min": best["fit_s"] if best else free_all["fit_s"],
+    }
+
+
+def main() -> int:
+    p = argparse.ArgumentParser()
+    p.add_argument("--backend", choices=["tabicl", "tabpfn"], required=True)
+    p.add_argument("--device", default="cuda:0")
+    p.add_argument("--task", default="classification")
+    p.add_argument("--cache-mode", default="kv")
+    p.add_argument("--shapes", nargs="+", type=parse_shape, required=True)
+    p.add_argument("--free-gib", nargs="+", type=float, default=[0, 24, 16, 8, 4, 2])
+    p.add_argument(
+        "--search", action="store_true", help="binary-search min free memory"
+    )
+    p.add_argument("--precision-gib", type=float, default=0.5)
+    p.add_argument("--out", default=None)
+    args = p.parse_args()
+
+    backend = build_backend(args.backend, args.device, args.cache_mode)
+    backend.warmup()
+    out = open(args.out, "a") if args.out else None  # noqa: SIM115
+    for rows, features in args.shapes:
+        records = (
+            [search(args, backend, rows, features)]
+            if args.search
+            else [run(args, backend, rows, features, free) for free in args.free_gib]
+        )
+        for record in records:
+            line = json.dumps(record)
+            print(line, flush=True)
+            if out:
+                out.write(line + "\n")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/benchmarks/eval_kv_precision.py b/benchmarks/eval_kv_precision.py
new file mode 100644
index 0000000..986c23e
--- /dev/null
+++ b/benchmarks/eval_kv_precision.py
@@ -0,0 +1,141 @@
+#!/usr/bin/env python3
+"""Accuracy, memory and latency of TabICL KV-cache precisions (#34).
+
+For each dataset, fits TabICLBackend once per precision (native / fp16 /
+int8) on the same split and compares against native: prediction
+agreement, max and mean probability drift, the task metric (accuracy and
+ROC AUC, or R^2), resident context bytes, and warm predict latency (what
+decompression costs).
+
+Datasets: the ones bundled with scikit-learn (no download) plus larger
+synthetic tables, where KV size actually matters.
+
+Usage: eval_kv_precision.py --device cuda:0 --out kv-precision-a100.json
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+import time
+
+import numpy as np
+import torch
+from sklearn.datasets import (
+    load_breast_cancer,
+    load_diabetes,
+    load_digits,
+    load_wine,
+    make_classification,
+    make_regression,
+)
+from sklearn.metrics import accuracy_score, r2_score, roc_auc_score
+from sklearn.model_selection import train_test_split
+
+from tabctx.backends.tabicl import TabICLBackend
+
+PRECISIONS = ("native", "fp16", "int8")
+
+
+def datasets():
+    yield "breast_cancer", "classification", *load_breast_cancer(return_X_y=True)
+    yield "digits(10cls)", "classification", *load_digits(return_X_y=True)
+    yield "wine(3cls)", "classification", *load_wine(return_X_y=True)
+    yield "diabetes", "regression", *load_diabetes(return_X_y=True)
+    X, y = make_classification(
+        20_000,
+        40,
+        n_informative=20,
+        n_classes=4,
+        n_clusters_per_class=2,
+        random_state=0,
+    )
+    yield "synth 20k x 40 (4cls)", "classification", X, y
+    X, y = make_classification(50_000, 10, n_informative=6, random_state=1)
+    yield "synth 50k x 10", "classification", X, y
+    X, y = make_regression(20_000, 30, n_informative=15, noise=5.0, random_state=2)
+    yield "synth 20k x 30 (reg)", "regression", X, y
+
+
+def metric(task: str, y_true, outcome) -> dict:
+    if task == "regression":
+        return {"r2": round(float(r2_score(y_true, outcome.predictions)), 5)}
+    pred = np.array(outcome.predictions)
+    out = {
+        "accuracy": round(float(accuracy_score(np.array(y_true).astype(str), pred)), 5)
+    }
+    proba = np.array(outcome.probabilities)
+    classes = outcome.classes
+    y_idx = np.array([classes.index(str(v)) for v in y_true])
+    if proba.shape[1] == 2:
+        out["auc"] = round(float(roc_auc_score(y_idx, proba[:, 1])), 5)
+    else:
+        out["auc"] = round(float(roc_auc_score(y_idx, proba, multi_class="ovr")), 5)
+    return out
+
+
+def main() -> int:
+    p = argparse.ArgumentParser()
+    p.add_argument("--device", default="cuda:0")
+    p.add_argument("--out", default=None)
+    args = p.parse_args()
+
+    backends = {
+        prec: TabICLBackend(device=args.device, kv_precision=prec)
+        for prec in PRECISIONS
+    }
+    for b in backends.values():
+        b.warmup()
+    results = []
+    for name, task, X, y in datasets():
+        X = np.asarray(X, dtype=np.float32)
+        X_tr, X_te, y_tr, y_te = train_test_split(X, y, test_size=0.25, random_state=0)
+        row = {
+            "dataset": name,
+            "task": task,
+            "train_rows": len(X_tr),
+            "features": X.shape[1],
+        }
+        base = None
+        for prec, b in backends.items():
+            payload = b.fit(X_tr, y_tr, task)
+            resident = b.context_bytes_hint(len(X_tr), X.shape[1]) or 0
+            out = b.predict(payload, X_te, return_proba=task == "classification")
+            torch.cuda.synchronize()
+            t0 = time.perf_counter()
+            for _ in range(5):
+                b.predict(payload, X_te[:100], False)
+            torch.cuda.synchronize()
+            entry = {
+                "resident_mib": round(resident / 2**20, 1),
+                "predict100_ms": round((time.perf_counter() - t0) / 5 * 1000, 1),
+                **metric(task, y_te, out),
+            }
+            if base is None:
+                base = out
+            elif task == "classification":
+                pa, pb = np.array(base.probabilities), np.array(out.probabilities)
+                entry["label_agreement"] = float(
+                    np.mean(np.array(base.predictions) == np.array(out.predictions))
+                )
+                entry["max_prob_diff"] = float(np.abs(pa - pb).max())
+                entry["mean_prob_diff"] = float(np.abs(pa - pb).mean())
+            else:
+                a, c = np.array(base.predictions), np.array(out.predictions)
+                entry["max_rel_diff"] = float(
+                    np.abs(a - c).max() / (np.abs(a).max() + 1e-9)
+                )
+            row[prec] = entry
+            del payload
+            torch.cuda.empty_cache()
+        results.append(row)
+        print(json.dumps(row), flush=True)
+    if args.out:
+        with open(args.out, "w") as f:
+            json.dump(results, f, indent=2)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/benchmarks/profile_fit_predict.py b/benchmarks/profile_fit_predict.py
index 426085e..52a54a5 100644
--- a/benchmarks/profile_fit_predict.py
+++ b/benchmarks/profile_fit_predict.py
@@ -89,42 +89,55 @@ def stage_modules(kind: str, backend, payload, task: str) -> dict[str, torch.nn.
 
 
 class StageRecorder:
-    """CUDA-event timing and peak memory per stage via forward hooks."""
+    """CUDA-event timing and peak memory per stage.
+
+    Wraps each stage's entry methods (forward, plus TabICL's cache-mode
+    entry points forward_with_cache / prepare_repr_cache /
+    forward_with_repr_cache, which bypass nn.Module hooks), recording only
+    the outermost call so a method calling another isn't counted twice."""
+
+    ENTRY_METHODS = (
+        "forward",
+        "forward_with_cache",
+        "prepare_repr_cache",
+        "forward_with_repr_cache",
+    )
 
     def __init__(self, device: str, modules: dict[str, torch.nn.Module]) -> None:
         self.device = torch.device(device)
         self.calls: dict[str, int] = defaultdict(int)
         self.time_ms: dict[str, float] = defaultdict(float)
         self.peak_bytes: dict[str, int] = defaultdict(int)
-        self._open: dict[str, tuple] = {}
-        self._handles = []
+        self._depth: dict[str, int] = defaultdict(int)
+        self._patched: list[tuple[torch.nn.Module, str]] = []
         for name, module in modules.items():
-            self._handles.append(module.register_forward_pre_hook(self._pre(name)))
-            self._handles.append(module.register_forward_hook(self._post(name)))
-
-    def _pre(self, name):
-        def hook(module, args):
+            for method in self.ENTRY_METHODS:
+                if callable(getattr(module, method, None)):
+                    setattr(module, method, self._wrap(name, getattr(module, method)))
+                    self._patched.append((module, method))
+
+    def _wrap(self, name, fn):
+        def wrapped(*args, **kwargs):
+            if self._depth[name]:
+                return fn(*args, **kwargs)
+            self._depth[name] += 1
             torch.cuda.synchronize(self.device)
             base = torch.cuda.memory_allocated(self.device)
             torch.cuda.reset_peak_memory_stats(self.device)
-            start = torch.cuda.Event(enable_timing=True)
+            start, end = (torch.cuda.Event(enable_timing=True) for _ in range(2))
             start.record()
-            self._open[name] = (start, base)
-
-        return hook
-
-    def _post(self, name):
-        def hook(module, args, output):
-            start, base = self._open.pop(name)
-            end = torch.cuda.Event(enable_timing=True)
-            end.record()
-            torch.cuda.synchronize(self.device)
-            self.calls[name] += 1
-            self.time_ms[name] += start.elapsed_time(end)
-            peak = torch.cuda.max_memory_allocated(self.device) - base
-            self.peak_bytes[name] = max(self.peak_bytes[name], peak)
-
-        return hook
+            try:
+                return fn(*args, **kwargs)
+            finally:
+                end.record()
+                torch.cuda.synchronize(self.device)
+                self._depth[name] -= 1
+                self.calls[name] += 1
+                self.time_ms[name] += start.elapsed_time(end)
+                peak = torch.cuda.max_memory_allocated(self.device) - base
+                self.peak_bytes[name] = max(self.peak_bytes[name], peak)
+
+        return wrapped
 
     def reset(self) -> None:
         self.calls.clear()
@@ -142,8 +155,8 @@ def summary(self) -> dict:
         }
 
     def close(self) -> None:
-        for h in self._handles:
-            h.remove()
+        for module, method in self._patched:
+            delattr(module, method)  # the instance attribute shadowing the method
 
 
 def timed(device: torch.device, fn):
@@ -260,7 +273,7 @@ def main() -> int:
 
     backend = build_backend(args.backend, args.device, args.cache_mode)
     warm_X, warm_y = make_table(200, 8, args.task, seed=0)
-    backend.fit(warm_X, warm_y, args.task)  # load weights outside measurements
+    backend.warmup()  # load weights outside measurements
     recorder_box: dict = {}
     out = open(args.out, "a") if args.out else None  # noqa: SIM115
     for rows, features in args.shapes:
diff --git a/src/tabctx/backends/kv_precision.py b/src/tabctx/backends/kv_precision.py
new file mode 100644
index 0000000..f24c78a
--- /dev/null
+++ b/src/tabctx/backends/kv_precision.py
@@ -0,0 +1,9 @@
+"""KV-cache storage precisions (torch-free, so settings can validate them
+without importing a GPU backend). See backends/tabicl_kv.py."""
+
+from __future__ import annotations
+
+from typing import Literal
+
+KVPrecision = Literal["native", "fp16", "int8"]
+KV_PRECISIONS: tuple[KVPrecision, ...] = ("native", "fp16", "int8")
diff --git a/src/tabctx/backends/tabicl.py b/src/tabctx/backends/tabicl.py
index 2329fb8..c170ad5 100644
--- a/src/tabctx/backends/tabicl.py
+++ b/src/tabctx/backends/tabicl.py
@@ -55,6 +55,7 @@
 import numpy as np
 
 from tabctx.backends._cuda import device_scope, ensure_initialized, is_cuda
+from tabctx.backends.tabicl_kv import KV_PRECISIONS, KVPrecision, compress_icl_cache
 from tabctx.errors import BackendComputeError
 from tabctx.types import ArrayLike, PredictOutcome, Task
 
@@ -90,7 +91,12 @@ class TabICLBackend:
     # named for the checkpoint actually loaded (tabicl-classifier-v2-*).
     name = "tabicl-v2"
 
-    def __init__(self, device: str | None = None, kv_cache: bool | str = "kv") -> None:
+    def __init__(
+        self,
+        device: str | None = None,
+        kv_cache: bool | str = "kv",
+        kv_precision: KVPrecision = "native",
+    ) -> None:
         """kv_cache: passed through to TabICL ("kv", "repr", True, or
         False). tabicl's own default is False, which makes every
         predict() re-encode the ENTIRE training set through all three
@@ -107,6 +113,10 @@ def __init__(self, device: str | None = None, kv_cache: bool | str = "kv") -> No
         self._device = device or ("cuda" if torch.cuda.is_available() else "cpu")
         ensure_initialized(self._device)
         self._kv_cache = kv_cache
+        if kv_precision not in KV_PRECISIONS:
+            raise ValueError(f"kv_precision must be one of {KV_PRECISIONS}")
+        # Storage of the ICL key/value cache (see backends/tabicl_kv.py).
+        self._kv_precision: KVPrecision = kv_precision
         self._last_context_bytes: int | None = None
         self._last_fit_peak_bytes: int | None = None
         # task -> attrs of an already-loaded backbone, shared across
@@ -173,6 +183,9 @@ def fit(self, X: ArrayLike, y: ArrayLike, task: Task) -> Any:
         try:
             with device_scope(device):
                 model.fit(X, y)
+                # Before measuring: the context's resident size is the
+                # compressed one.
+                compress_icl_cache(model, self._kv_precision)
         except torch.cuda.OutOfMemoryError as e:
             self._last_context_bytes = None
             self._last_fit_peak_bytes = None
diff --git a/src/tabctx/backends/tabicl_kv.py b/src/tabctx/backends/tabicl_kv.py
new file mode 100644
index 0000000..2fe46a5
--- /dev/null
+++ b/src/tabctx/backends/tabicl_kv.py
@@ -0,0 +1,152 @@
+"""Compressed storage for TabICL's ICL key/value cache.
+
+A TabICL context in kv mode is dominated by the ICL transformer's cached
+keys and values: ~0.19 MiB per training row for the default 8 estimators
+(measured on an A100: 100k x 10 = 17.6 GiB resident), stored in fp16 once
+AMP is on (>= 1024 rows) and fp32 below that. That is what bounds the
+longest table a GPU can hold, and how many contexts fit.
+
+This module swaps the ICL cache's entries for ``CompressedKVEntry``, which
+stores them compactly and decompresses on access:
+
+- ``"int8"``: symmetric int8 per (head, row) over the head dimension, with
+  an fp16 scale -- about 0.52x of fp16 (1 byte + 2/head_dim bytes per
+  element), 0.26x of fp32.
+- ``"fp16"``: half-precision storage of fp32 caches (small tables), 0.5x;
+  a no-op for caches already in fp16.
+
+Attention reads ``entry.key`` / ``entry.value`` once per layer per predict
+(tabicl/_model/attention.py), so only one layer's K/V is ever materialized
+at full precision, transiently. Decompression returns the original dtype,
+so the model's compute path is unchanged.
+
+``CompressedKVEntry`` duck-types tabicl's ``KVCacheEntry`` (is_valid,
+batch slicing via __getitem__/__setitem__, ``to``), which is everything
+tabicl's InferenceManager batching and its pickling (``cache.to("cpu")``
+on save, ``.to(device)`` on load) use.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+from torch import Tensor
+
+from tabctx.backends.kv_precision import KV_PRECISIONS, KVPrecision
+
+__all__ = ["KV_PRECISIONS", "CompressedKVEntry", "KVPrecision", "compress_icl_cache"]
+
+
+class _Packed:
+    """One compressed tensor: payload (+ scale for int8) and the dtype to
+    decompress to."""
+
+    __slots__ = ("data", "scale", "dtype")
+
+    def __init__(self, data: Tensor, scale: Tensor | None, dtype: torch.dtype) -> None:
+        self.data = data
+        self.scale = scale
+        self.dtype = dtype
+
+    @classmethod
+    def pack(cls, t: Tensor, precision: KVPrecision) -> _Packed:
+        if precision == "int8":
+            absmax = t.abs().amax(dim=-1, keepdim=True).float().clamp_min(1e-8)
+            scale = (absmax / 127.0).to(torch.float16)
+            q = torch.round(t.float() / scale.float()).clamp_(-127, 127).to(torch.int8)
+            return cls(q, scale, t.dtype)
+        return cls(t.to(torch.float16), None, t.dtype)
+
+    def unpack(self) -> Tensor:
+        if self.scale is None:
+            return self.data.to(self.dtype)
+        return (self.data.to(torch.float16) * self.scale).to(self.dtype)
+
+    def slice(self, indices: Any) -> _Packed:
+        # Scale has the same leading dims as data (last dim 1).
+        scale = None if self.scale is None else self.scale[indices]
+        return _Packed(self.data[indices], scale, self.dtype)
+
+    def write(self, indices: Any, other: _Packed) -> None:
+        self.data[indices] = other.data
+        if self.scale is not None and other.scale is not None:
+            self.scale[indices] = other.scale
+
+    def to(self, device: Any) -> _Packed:
+        scale = None if self.scale is None else self.scale.to(device)
+        return _Packed(self.data.to(device), scale, self.dtype)
+
+    @property
+    def nbytes(self) -> int:
+        n = self.data.numel() * self.data.element_size()
+        if self.scale is not None:
+            n += self.scale.numel() * self.scale.element_size()
+        return n
+
+
+class CompressedKVEntry:
+    """A KVCacheEntry stored compressed; ``key``/``value`` decompress."""
+
+    def __init__(self, key: _Packed, value: _Packed) -> None:
+        self._key = key
+        self._value = value
+
+    @classmethod
+    def compress(cls, entry: Any, precision: KVPrecision) -> CompressedKVEntry:
+        return cls(
+            _Packed.pack(entry.key, precision), _Packed.pack(entry.value, precision)
+        )
+
+    @property
+    def key(self) -> Tensor:
+        return self._key.unpack()
+
+    @property
+    def value(self) -> Tensor:
+        return self._value.unpack()
+
+    @property
+    def nbytes(self) -> int:
+        return self._key.nbytes + self._value.nbytes
+
+    def is_valid(self) -> bool:
+        return True
+
+    def __getitem__(self, indices: Any) -> CompressedKVEntry:
+        return CompressedKVEntry(self._key.slice(indices), self._value.slice(indices))
+
+    def __setitem__(self, indices: Any, other: Any) -> None:
+        if isinstance(other, CompressedKVEntry):
+            self._key.write(indices, other._key)
+            self._value.write(indices, other._value)
+        else:  # a plain KVCacheEntry: compress into our layout
+            precision: KVPrecision = "int8" if self._key.scale is not None else "fp16"
+            self[indices] = CompressedKVEntry.compress(other, precision)
+
+    def to(self, device: Any, dtype: Any = None) -> CompressedKVEntry:
+        del dtype  # storage dtype is fixed by the compression mode
+        return CompressedKVEntry(self._key.to(device), self._value.to(device))
+
+
+def compress_icl_cache(estimator: Any, precision: KVPrecision) -> int:
+    """Compress the ICL K/V of every cache on a fitted TabICL estimator, in
+    place. Returns the number of entries compressed (0 for "native", for
+    estimators without a kv cache, or when already compressed)."""
+    if precision == "native":
+        return 0
+    caches = getattr(estimator, "model_kv_cache_", None) or {}
+    n = 0
+    for cache in caches.values():
+        icl = getattr(cache, "icl_cache", None)
+        kv = getattr(icl, "kv", None)
+        if not kv:
+            continue
+        for idx, entry in list(kv.items()):
+            if isinstance(entry, CompressedKVEntry) or not entry.is_valid():
+                continue
+            if precision == "fp16" and entry.key.dtype == torch.float16:
+                continue  # already half precision
+            kv[idx] = CompressedKVEntry.compress(entry, precision)
+            n += 1
+    return n
diff --git a/src/tabctx/serve/factory.py b/src/tabctx/serve/factory.py
index 2ee2e36..bb9e8d1 100644
--- a/src/tabctx/serve/factory.py
+++ b/src/tabctx/serve/factory.py
@@ -28,6 +28,11 @@
   fastest per predict; "repr" uses far less cache memory per context;
   "off" re-encodes the training set on every predict (tabicl's own
   default, kept only as an escape hatch).
+- ``TABCTX_KV_PRECISION``: ``"native"`` (default), ``"fp16"`` or ``"int8"``
+  -- storage of TabICL's ICL key/value cache, which dominates a TabICL
+  context (see backends/tabicl_kv.py). ``"int8"`` about halves resident
+  memory per context (and so doubles the longest table and the contexts
+  per GPU). TabPFN already stores its cache in int8 by default.
 - ``TABCTX_BATCH_WINDOW_MS``: coalescing window for same-context predict
   batching (see batching.py); default 5. 0 disables coalescing.
 - ``TABCTX_MAX_UPLOAD_BYTES`` (default 4GiB) and ``TABCTX_UPLOAD_TTL_S``
@@ -58,6 +63,7 @@
 from typing import Literal
 
 from tabctx.backends.base import TabularICLBackend
+from tabctx.backends.kv_precision import KV_PRECISIONS, KVPrecision
 from tabctx.cache.manager import ContextCacheManager
 from tabctx.cache.spill import DEFAULT_SPILL_CAPACITY_BYTES, DiskSpillStore
 from tabctx.engine import TabctxEngine
@@ -77,6 +83,7 @@
 BACKEND_ENV_VAR = "TABCTX_BACKEND"
 GPU_MEMORY_FRACTION_ENV_VAR = "TABCTX_GPU_MEMORY_FRACTION"
 KV_CACHE_ENV_VAR = "TABCTX_KV_CACHE"
+KV_PRECISION_ENV_VAR = "TABCTX_KV_PRECISION"
 BATCH_WINDOW_MS_ENV_VAR = "TABCTX_BATCH_WINDOW_MS"
 MAX_UPLOAD_BYTES_ENV_VAR = "TABCTX_MAX_UPLOAD_BYTES"
 SPILL_DIR_ENV_VAR = "TABCTX_SPILL_DIR"
@@ -97,6 +104,7 @@ class ServeSettings:
     backends: tuple[BackendKind, ...] = ("tabicl",)
     gpu_memory_fraction: float = 1.0
     kv_cache: KvCacheMode = "kv"
+    kv_precision: KVPrecision = "native"
     batch_window_ms: float = 5.0
     max_upload_bytes: int = DEFAULT_MAX_UPLOAD_BYTES
     upload_ttl_s: float = DEFAULT_TTL_S
@@ -135,6 +143,14 @@ def from_env(cls) -> ServeSettings:
                 f"{KV_CACHE_ENV_VAR}={kv_cache!r} is not a known mode "
                 "(expected 'kv', 'repr', or 'off')"
             )
+        kv_precision = (
+            os.environ.get(KV_PRECISION_ENV_VAR, cls.kv_precision).strip().lower()
+        )
+        if kv_precision not in KV_PRECISIONS:
+            raise ValueError(
+                f"{KV_PRECISION_ENV_VAR}={kv_precision!r} must be one of "
+                f"{', '.join(KV_PRECISIONS)}"
+            )
         raw_window = os.environ.get(BATCH_WINDOW_MS_ENV_VAR, str(cls.batch_window_ms))
         try:
             batch_window_ms = float(raw_window)
@@ -175,6 +191,7 @@ def from_env(cls) -> ServeSettings:
             backends=backends,
             gpu_memory_fraction=fraction,
             kv_cache=kv_cache,
+            kv_precision=kv_precision,
             batch_window_ms=batch_window_ms,
             max_upload_bytes=max_upload_bytes,
             upload_ttl_s=upload_ttl_s,
@@ -377,7 +394,9 @@ def _build_real_backend(
         kv_cache: bool | str = (
             False if settings.kv_cache == "off" else settings.kv_cache
         )
-        backend = TabICLBackend(device=device, kv_cache=kv_cache)
+        backend = TabICLBackend(
+            device=device, kv_cache=kv_cache, kv_precision=settings.kv_precision
+        )
     # Weights load at replica start, not inside (and charged to) the first
     # tenant's fit.
     backend.warmup()
diff --git a/tests/unit/test_serve_factory.py b/tests/unit/test_serve_factory.py
index 2417d23..258b26a 100644
--- a/tests/unit/test_serve_factory.py
+++ b/tests/unit/test_serve_factory.py
@@ -252,6 +252,7 @@ def test_settings_defaults_match_from_env_defaults(self, monkeypatch):
             SPILL_DIR_ENV_VAR,
             SPILL_CAPACITY_ENV_VAR,
             DEVICES_ENV_VAR,
+            "TABCTX_KV_PRECISION",
         ):
             monkeypatch.delenv(var, raising=False)
         assert ServeSettings.from_env() == ServeSettings()
@@ -451,3 +452,14 @@ def test_spill_layout_change_is_reported(self, tmp_path, caplog):
                 )
             )
         assert "will not be restored" in caplog.text
+
+
+class TestKvPrecision:
+    def test_default_and_parsing(self, monkeypatch):
+        monkeypatch.delenv("TABCTX_KV_PRECISION", raising=False)
+        assert ServeSettings.from_env().kv_precision == "native"
+        monkeypatch.setenv("TABCTX_KV_PRECISION", " INT8 ")
+        assert ServeSettings.from_env().kv_precision == "int8"
+        monkeypatch.setenv("TABCTX_KV_PRECISION", "fp4")
+        with pytest.raises(ValueError, match="TABCTX_KV_PRECISION"):
+            ServeSettings.from_env()
diff --git a/tests/unit/test_tabicl_kv.py b/tests/unit/test_tabicl_kv.py
new file mode 100644
index 0000000..fd9dcc1
--- /dev/null
+++ b/tests/unit/test_tabicl_kv.py
@@ -0,0 +1,107 @@
+"""Unit tests for compressed TabICL ICL K/V storage (no tabicl needed)."""
+
+import pickle
+from dataclasses import dataclass
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+from tabctx.backends.tabicl_kv import CompressedKVEntry, compress_icl_cache  # noqa: E402
+
+
+@dataclass
+class Entry:  # the shape of tabicl's KVCacheEntry
+    key: object = None
+    value: object = None
+
+    def is_valid(self):
+        return self.key is not None and self.value is not None
+
+
+def entry(dtype=torch.float16, shape=(4, 8, 300, 64), seed=0):
+    g = torch.Generator().manual_seed(seed)
+    k = (torch.randn(shape, generator=g) * 3).to(dtype)
+    v = torch.randn(shape, generator=g).to(dtype)
+    return Entry(k, v)
+
+
+def nbytes(t):
+    return t.numel() * t.element_size()
+
+
+class TestCompressedKVEntry:
+    @pytest.mark.parametrize("dtype", [torch.float16, torch.float32])
+    def test_int8_roundtrip_is_close_and_keeps_dtype(self, dtype):
+        e = entry(dtype)
+        c = CompressedKVEntry.compress(e, "int8")
+        for orig, back in ((e.key, c.key), (e.value, c.value)):
+            assert back.dtype == dtype and back.shape == orig.shape
+            # per-(head,row) absmax/127 quantization: error is half a step,
+            # plus fp16 rounding of the scale and the dequantize -- under one
+            # step (measured worst case: 0.54 steps, 0.4% of absmax)
+            err = (back.float() - orig.float()).abs()
+            bound = orig.float().abs().amax(-1, keepdim=True) / 127
+            assert bool((err <= bound).all())
+
+    def test_int8_is_about_half_of_fp16(self):
+        e = entry(torch.float16)
+        c = CompressedKVEntry.compress(e, "int8")
+        ratio = c.nbytes / (nbytes(e.key) + nbytes(e.value))
+        assert 0.5 < ratio < 0.53  # 1 byte + 2/64 bytes vs 2 bytes
+
+    def test_fp16_halves_fp32_exactly_representable(self):
+        e = entry(torch.float32)
+        c = CompressedKVEntry.compress(e, "fp16")
+        assert c.nbytes * 2 == nbytes(e.key) + nbytes(e.value)
+        assert c.key.dtype == torch.float32
+        assert torch.equal(c.key, e.key.half().float())
+
+    def test_batch_slicing_matches_full_unpack(self):
+        c = CompressedKVEntry.compress(entry(), "int8")
+        assert torch.equal(c[1:3].key, c.key[1:3])
+        assert torch.equal(c[(slice(0, 2), slice(None))].value, c.value[0:2])
+
+    def test_setitem_writes_a_batch_slice(self):
+        c = CompressedKVEntry.compress(entry(seed=0), "int8")
+        other = CompressedKVEntry.compress(entry(shape=(2, 8, 300, 64), seed=1), "int8")
+        c[0:2] = other
+        assert torch.equal(c.key[0:2], other.key)
+        plain = entry(shape=(1, 8, 300, 64), seed=2)
+        c[3:4] = plain  # a plain entry is compressed into the layout
+        assert torch.allclose(c.key[3:4].float(), plain.key.float(), atol=0.1)
+
+    def test_to_and_pickle_roundtrip(self):
+        c = CompressedKVEntry.compress(entry(), "int8")
+        moved = c.to("cpu", dtype=torch.float32)  # dtype is ignored: storage is fixed
+        assert torch.equal(moved.key, c.key)
+        restored = pickle.loads(pickle.dumps(c))
+        assert torch.equal(restored.value, c.value) and restored.is_valid()
+
+
+class _Cache:
+    def __init__(self, entries):
+        self.icl_cache = type("KV", (), {"kv": dict(enumerate(entries))})()
+
+
+class TestCompressIclCache:
+    def test_compresses_every_icl_entry_once(self):
+        est = type("Est", (), {})()
+        est.model_kv_cache_ = {
+            "power": _Cache([entry(), entry()]),
+            "none": _Cache([entry()]),
+        }
+        assert compress_icl_cache(est, "int8") == 3
+        assert compress_icl_cache(est, "int8") == 0  # idempotent
+        assert all(
+            isinstance(e, CompressedKVEntry)
+            for c in est.model_kv_cache_.values()
+            for e in c.icl_cache.kv.values()
+        )
+
+    def test_native_and_fp16_on_fp16_are_noops(self):
+        est = type("Est", (), {})()
+        est.model_kv_cache_ = {"m": _Cache([entry(torch.float16)])}
+        assert compress_icl_cache(est, "native") == 0
+        assert compress_icl_cache(est, "fp16") == 0
+        assert compress_icl_cache(type("NoCache", (), {})(), "int8") == 0

From 896550ae563c9fbcf41dff3d56f3866ac007d28c Mon Sep 17 00:00:00 2001
From: Amrit Krishnan <amrit110@gmail.com>
Date: Sat, 17 Oct 2026 15:45:39 -0400
Subject: [PATCH 02/39] perf(tabicl): int8 ICL KV cache by default (A100 study:
 identical labels/metrics on 7 datasets, ~0.52x resident)

Also: devbox uses a private kubeconfig (get-credentials switched the
current context under a concurrently running rehearsal).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../baselines/kv-precision-tabicl-a100.jsonl     |  7 +++++++
 .../baselines/minfree-tabicl-kv-a100.jsonl       | 16 ++++++++++++++++
 benchmarks/gke/devbox.sh                         |  3 +++
 src/tabctx/serve/factory.py                      | 11 +++++++----
 tests/unit/test_serve_factory.py                 |  2 +-
 5 files changed, 34 insertions(+), 5 deletions(-)
 create mode 100644 benchmarks/baselines/kv-precision-tabicl-a100.jsonl
 create mode 100644 benchmarks/baselines/minfree-tabicl-kv-a100.jsonl

diff --git a/benchmarks/baselines/kv-precision-tabicl-a100.jsonl b/benchmarks/baselines/kv-precision-tabicl-a100.jsonl
new file mode 100644
index 0000000..20ffa88
--- /dev/null
+++ b/benchmarks/baselines/kv-precision-tabicl-a100.jsonl
@@ -0,0 +1,7 @@
+{"dataset": "breast_cancer", "task": "classification", "train_rows": 426, "features": 30, "native": {"resident_mib": 262.7, "predict100_ms": 58.2, "accuracy": 0.97902, "auc": 0.99916}, "fp16": {"resident_mib": 182.3, "predict100_ms": 55.9, "accuracy": 0.97902, "auc": 0.99916, "label_agreement": 1.0, "max_prob_diff": 2.6404857635498047e-05, "mean_prob_diff": 1.5774988629867426e-06}, "int8": {"resident_mib": 143.2, "predict100_ms": 64.3, "accuracy": 0.97902, "auc": 0.99916, "label_agreement": 1.0, "max_prob_diff": 0.0009291172027587891, "mean_prob_diff": 5.036376046995861e-05}}
+{"dataset": "digits(10cls)", "task": "classification", "train_rows": 1347, "features": 64, "native": {"resident_mib": 353.1, "predict100_ms": 88.1, "accuracy": 0.99778, "auc": 0.99996}, "fp16": {"resident_mib": 353.1, "predict100_ms": 84.6, "accuracy": 0.99778, "auc": 0.99996, "label_agreement": 1.0, "max_prob_diff": 0.0, "mean_prob_diff": 0.0}, "int8": {"resident_mib": 228.4, "predict100_ms": 86.6, "accuracy": 0.99778, "auc": 0.99996, "label_agreement": 1.0, "max_prob_diff": 0.0024799108505249023, "mean_prob_diff": 3.1142769239660417e-06}}
+{"dataset": "wine(3cls)", "task": "classification", "train_rows": 133, "features": 13, "native": {"resident_mib": 105.2, "predict100_ms": 50.3, "accuracy": 1.0, "auc": 1.0}, "fp16": {"resident_mib": 76.9, "predict100_ms": 46.3, "accuracy": 1.0, "auc": 1.0, "label_agreement": 1.0, "max_prob_diff": 7.152557373046875e-06, "mean_prob_diff": 1.7900499516245974e-07}, "int8": {"resident_mib": 64.9, "predict100_ms": 50.4, "accuracy": 1.0, "auc": 1.0, "label_agreement": 1.0, "max_prob_diff": 0.00021535158157348633, "mean_prob_diff": 8.665528010051613e-06}}
+{"dataset": "diabetes", "task": "regression", "train_rows": 331, "features": 10, "native": {"resident_mib": 170.3, "predict100_ms": 50.8, "r2": 0.35805}, "fp16": {"resident_mib": 114.2, "predict100_ms": 52.0, "r2": 0.35806, "max_rel_diff": 1.1122660013609997e-05}, "int8": {"resident_mib": 77.1, "predict100_ms": 53.6, "r2": 0.35803, "max_rel_diff": 0.00014799947609946365}}
+{"dataset": "synth 20k x 40 (4cls)", "task": "classification", "train_rows": 15000, "features": 40, "native": {"resident_mib": 2885.7, "predict100_ms": 83.6, "accuracy": 0.9826, "auc": 0.99582}, "fp16": {"resident_mib": 2885.7, "predict100_ms": 79.1, "accuracy": 0.9826, "auc": 0.99582, "label_agreement": 1.0, "max_prob_diff": 0.0, "mean_prob_diff": 0.0}, "int8": {"resident_mib": 1520.9, "predict100_ms": 95.5, "accuracy": 0.9826, "auc": 0.99582, "label_agreement": 1.0, "max_prob_diff": 0.0032877326011657715, "mean_prob_diff": 2.0953629555515362e-05}}
+{"dataset": "synth 50k x 10", "task": "classification", "train_rows": 37500, "features": 10, "native": {"resident_mib": 7052.9, "predict100_ms": 80.4, "accuracy": 0.992, "auc": 0.995}, "fp16": {"resident_mib": 7052.9, "predict100_ms": 79.9, "accuracy": 0.992, "auc": 0.995, "label_agreement": 1.0, "max_prob_diff": 0.0, "mean_prob_diff": 0.0}, "int8": {"resident_mib": 3646.5, "predict100_ms": 110.3, "accuracy": 0.992, "auc": 0.995, "label_agreement": 1.0, "max_prob_diff": 0.00897100567817688, "mean_prob_diff": 1.816937566152774e-05}}
+{"dataset": "synth 20k x 30 (reg)", "task": "regression", "train_rows": 15000, "features": 30, "native": {"resident_mib": 2872.6, "predict100_ms": 81.3, "r2": 0.99921}, "fp16": {"resident_mib": 2872.6, "predict100_ms": 94.8, "r2": 0.99921, "max_rel_diff": 0.0}, "int8": {"resident_mib": 1507.9, "predict100_ms": 91.4, "r2": 0.99921, "max_rel_diff": 0.0003313215243717279}}
diff --git a/benchmarks/baselines/minfree-tabicl-kv-a100.jsonl b/benchmarks/baselines/minfree-tabicl-kv-a100.jsonl
new file mode 100644
index 0000000..11cdfaa
--- /dev/null
+++ b/benchmarks/baselines/minfree-tabicl-kv-a100.jsonl
@@ -0,0 +1,16 @@
+{"rows": 2000, "features": 10, "status": "ok", "resident_gib": 0.39, "unconstrained_peak_gib": 0.59, "unconstrained_fit_s": 0.433, "min_free_gib": 0.69, "min_working_gib": 0.3, "fit_s_at_min": 0.277}
+{"rows": 20000, "features": 10, "status": "ok", "resident_gib": 3.68, "unconstrained_peak_gib": 5.53, "unconstrained_fit_s": 2.141, "min_free_gib": 5.32, "min_working_gib": 1.64, "fit_s_at_min": 1.327}
+{"rows": 45000, "features": 10, "status": "ok", "resident_gib": 6.58, "unconstrained_peak_gib": 10.72, "unconstrained_fit_s": 4.581, "min_free_gib": 14.62, "min_working_gib": 8.04, "fit_s_at_min": 4.051}
+{"rows": 80000, "features": 10, "status": "ok", "resident_gib": 11.23, "unconstrained_peak_gib": 18.59, "unconstrained_fit_s": 10.753, "min_free_gib": 20.69, "min_working_gib": 9.46, "fit_s_at_min": 9.878}
+{"rows": 2000, "features": 50, "status": "ok", "resident_gib": 0.45, "unconstrained_peak_gib": 1.64, "unconstrained_fit_s": 0.694, "min_free_gib": 1.05, "min_working_gib": 0.6, "fit_s_at_min": 0.74}
+{"rows": 10000, "features": 50, "status": "ok", "resident_gib": 1.92, "unconstrained_peak_gib": 7.76, "unconstrained_fit_s": 1.957, "min_free_gib": 2.78, "min_working_gib": 0.86, "fit_s_at_min": 1.764}
+{"rows": 40000, "features": 50, "status": "ok", "resident_gib": 7.4, "unconstrained_peak_gib": 25.85, "unconstrained_fit_s": 6.619, "min_free_gib": 11.8, "min_working_gib": 4.4, "fit_s_at_min": 8.401}
+{"rows": 2000, "features": 100, "status": "ok", "resident_gib": 0.52, "unconstrained_peak_gib": 2.98, "unconstrained_fit_s": 1.273, "min_free_gib": 1.12, "min_working_gib": 0.6, "fit_s_at_min": 1.707}
+{"rows": 20000, "features": 100, "status": "ok", "resident_gib": 3.81, "unconstrained_peak_gib": 24.94, "unconstrained_fit_s": 5.871, "min_free_gib": 5.45, "min_working_gib": 1.64, "fit_s_at_min": 8.606}
+{"rows": 50000, "features": 100, "status": "ok", "resident_gib": 7.72, "unconstrained_peak_gib": 30.15, "unconstrained_fit_s": 13.99, "min_free_gib": 13.04, "min_working_gib": 5.32, "fit_s_at_min": 21.466}
+{"rows": 1000, "features": 300, "status": "ok", "resident_gib": 0.63, "unconstrained_peak_gib": 4.48, "unconstrained_fit_s": 3.021, "min_free_gib": 1.22, "min_working_gib": 0.59, "fit_s_at_min": 3.665}
+{"rows": 5000, "features": 300, "status": "ok", "resident_gib": 1.38, "unconstrained_peak_gib": 20.0, "unconstrained_fit_s": 5.614, "min_free_gib": 2.25, "min_working_gib": 0.87, "fit_s_at_min": 8.485}
+{"rows": 1000, "features": 500, "status": "ok", "resident_gib": 0.92, "unconstrained_peak_gib": 7.36, "unconstrained_fit_s": 5.105, "min_free_gib": 1.51, "min_working_gib": 0.59, "fit_s_at_min": 6.22}
+{"rows": 5000, "features": 500, "status": "ok", "resident_gib": 1.68, "unconstrained_peak_gib": 25.27, "unconstrained_fit_s": 9.848, "min_free_gib": 2.84, "min_working_gib": 1.16, "fit_s_at_min": 13.435}
+{"rows": 500, "features": 1000, "status": "ok", "resident_gib": 1.56, "unconstrained_peak_gib": 8.26, "unconstrained_fit_s": 9.29, "min_free_gib": 2.43, "min_working_gib": 0.87, "fit_s_at_min": 9.966}
+{"rows": 2000, "features": 1000, "status": "ok", "resident_gib": 1.84, "unconstrained_peak_gib": 24.56, "unconstrained_fit_s": 13.093, "min_free_gib": 2.99, "min_working_gib": 1.15, "fit_s_at_min": 16.456}
diff --git a/benchmarks/gke/devbox.sh b/benchmarks/gke/devbox.sh
index 4e71960..84d7a61 100755
--- a/benchmarks/gke/devbox.sh
+++ b/benchmarks/gke/devbox.sh
@@ -29,6 +29,9 @@ POD=tabctx-devbox
 HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 REPO="$(cd "$HERE/../.." && pwd)"
 mkdir -p "$STATE"
+# A private kubeconfig: get-credentials would otherwise switch the user's (and any
+# concurrently running script's) current kubectl context to this cluster.
+export KUBECONFIG="$STATE/kubeconfig"
 
 accelerator_for() {
   case "$1" in
diff --git a/src/tabctx/serve/factory.py b/src/tabctx/serve/factory.py
index bb9e8d1..b37994e 100644
--- a/src/tabctx/serve/factory.py
+++ b/src/tabctx/serve/factory.py
@@ -28,11 +28,14 @@
   fastest per predict; "repr" uses far less cache memory per context;
   "off" re-encodes the training set on every predict (tabicl's own
   default, kept only as an escape hatch).
-- ``TABCTX_KV_PRECISION``: ``"native"`` (default), ``"fp16"`` or ``"int8"``
+- ``TABCTX_KV_PRECISION``: ``"int8"`` (default), ``"fp16"`` or ``"native"``
   -- storage of TabICL's ICL key/value cache, which dominates a TabICL
   context (see backends/tabicl_kv.py). ``"int8"`` about halves resident
-  memory per context (and so doubles the longest table and the contexts
-  per GPU). TabPFN already stores its cache in int8 by default.
+  memory per context, so it doubles the longest table and the contexts per
+  GPU. Measured on an A100 over 7 datasets
+  (benchmarks/baselines/kv-precision-tabicl-a100.jsonl): identical labels
+  and metrics, max probability drift 0.009, at up to ~1.4x warm predict
+  latency on large contexts. TabPFN already stores its cache in int8.
 - ``TABCTX_BATCH_WINDOW_MS``: coalescing window for same-context predict
   batching (see batching.py); default 5. 0 disables coalescing.
 - ``TABCTX_MAX_UPLOAD_BYTES`` (default 4GiB) and ``TABCTX_UPLOAD_TTL_S``
@@ -104,7 +107,7 @@ class ServeSettings:
     backends: tuple[BackendKind, ...] = ("tabicl",)
     gpu_memory_fraction: float = 1.0
     kv_cache: KvCacheMode = "kv"
-    kv_precision: KVPrecision = "native"
+    kv_precision: KVPrecision = "int8"
     batch_window_ms: float = 5.0
     max_upload_bytes: int = DEFAULT_MAX_UPLOAD_BYTES
     upload_ttl_s: float = DEFAULT_TTL_S
diff --git a/tests/unit/test_serve_factory.py b/tests/unit/test_serve_factory.py
index 258b26a..4a07467 100644
--- a/tests/unit/test_serve_factory.py
+++ b/tests/unit/test_serve_factory.py
@@ -457,7 +457,7 @@ def test_spill_layout_change_is_reported(self, tmp_path, caplog):
 class TestKvPrecision:
     def test_default_and_parsing(self, monkeypatch):
         monkeypatch.delenv("TABCTX_KV_PRECISION", raising=False)
-        assert ServeSettings.from_env().kv_precision == "native"
+        assert ServeSettings.from_env().kv_precision == "int8"
         monkeypatch.setenv("TABCTX_KV_PRECISION", " INT8 ")
         assert ServeSettings.from_env().kv_precision == "int8"
         monkeypatch.setenv("TABCTX_KV_PRECISION", "fp4")

From 996bdd8891be93dcfc85961258b962c90a75d9e0 Mon Sep 17 00:00:00 2001
From: Amrit Krishnan <amrit110@gmail.com>
Date: Sat, 17 Oct 2026 15:48:46 -0400
Subject: [PATCH 03/39] feat(memory): admit fits by resident + working floor,
 not unconstrained peak

TabICL sizes stage batches from free GPU memory, so its unconstrained
peak fills the card. Measured on an A100: fits need resident + a small
working floor (5000x500: 2.8 GB vs 25 GB peak; 20000x100: 5.5 vs 25).
WorkingFloorModel (memory/requirement.py) bounds the floor from measured
points; AdaptiveMemoryEstimator.record_fit learns resident + floor when
given one (peak otherwise, e.g. TabPFN until calibrated). Generated
A100 TabICL kv requirement grid preloaded beside the peak grid.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 benchmarks/gen_requirement_grid.py            |  83 ++++++++++++
 src/tabctx/engine.py                          |  23 ++--
 src/tabctx/memory/adaptive.py                 |  30 ++++-
 .../calibration_tabicl_a100_requirement.py    |  81 ++++++++++++
 src/tabctx/memory/requirement.py              |  72 +++++++++++
 src/tabctx/serve/factory.py                   |  26 +++-
 tests/unit/test_requirement.py                | 122 ++++++++++++++++++
 tests/unit/test_serve_factory.py              |   5 +-
 8 files changed, 428 insertions(+), 14 deletions(-)
 create mode 100644 benchmarks/gen_requirement_grid.py
 create mode 100644 src/tabctx/memory/calibration_tabicl_a100_requirement.py
 create mode 100644 src/tabctx/memory/requirement.py
 create mode 100644 tests/unit/test_requirement.py

diff --git a/benchmarks/gen_requirement_grid.py b/benchmarks/gen_requirement_grid.py
new file mode 100644
index 0000000..b986c15
--- /dev/null
+++ b/benchmarks/gen_requirement_grid.py
@@ -0,0 +1,83 @@
+#!/usr/bin/env python3
+"""Generate a requirement calibration module from bench_memory_pressure.py
+--search output (JSON lines): measured working floors (FloorPoint) and the
+resident + floor requirement of each shape (Observation), for
+WorkingFloorModel and AdaptiveMemoryEstimator (see memory/requirement.py).
+
+Usage:
+  gen_requirement_grid.py benchmarks/baselines/minfree-tabicl-kv-a100.jsonl \
+      --prefix A100_40GB_TABICL_KV --source "A100-SXM4-40GB, tabicl 2.2.0, kv" \
+      --out src/tabctx/memory/calibration_tabicl_a100_requirement.py
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+
+GiB = 1024**3
+
+
+def main() -> int:
+    p = argparse.ArgumentParser()
+    p.add_argument("inputs", nargs="+")
+    p.add_argument("--prefix", required=True)
+    p.add_argument("--source", required=True)
+    p.add_argument("--out", required=True)
+    args = p.parse_args()
+    rows = []
+    for path in args.inputs:
+        with open(path) as f:
+            rows += [json.loads(line) for line in f if line.strip()]
+    ok = sorted(
+        (r for r in rows if r.get("status") == "ok" and "min_working_gib" in r),
+        key=lambda r: (
+            r["n_features"] if "n_features" in r else r["features"],
+            r["rows"],
+        ),
+    )
+    floors, needs = [], []
+    for r in ok:
+        n, d = r["rows"], r["features"]
+        working = int(r["min_working_gib"] * GiB)
+        need = int(r["min_free_gib"] * GiB)
+        note = (
+            f"resident {r['resident_gib']} GiB, unconstrained peak "
+            f"{r['unconstrained_peak_gib']} GiB, fit {r['unconstrained_fit_s']}s "
+            f"-> {r['fit_s_at_min']}s at the floor"
+        )
+        floors.append(
+            f"    FloorPoint(n_train={n}, n_features={d}, working_bytes={working}),  # {note}"
+        )
+        needs.append(
+            f"    Observation(n_train={n}, n_features={d}, real_bytes={need}),"
+        )
+    body = f'''"""GENERATED by benchmarks/gen_requirement_grid.py -- do not hand-edit.
+
+Measured requirements ({args.source}): for each shape, the least free GPU
+memory with which a fit succeeded (resident context + working floor), from
+benchmarks/bench_memory_pressure.py --search. See memory/requirement.py.
+"""
+
+from __future__ import annotations
+
+from tabctx.memory.adaptive import Observation
+from tabctx.memory.requirement import FloorPoint
+
+{args.prefix}_FLOOR_POINTS: tuple[FloorPoint, ...] = (
+{chr(10).join(floors)}
+)
+
+{args.prefix}_REQUIREMENT_GRID: tuple[Observation, ...] = (
+{chr(10).join(needs)}
+)
+'''
+    with open(args.out, "w") as f:
+        f.write(body)
+    print(f"wrote {len(ok)} shapes to {args.out}")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/src/tabctx/engine.py b/src/tabctx/engine.py
index 1186cf5..d4cff88 100644
--- a/src/tabctx/engine.py
+++ b/src/tabctx/engine.py
@@ -256,18 +256,21 @@ def fit(
             if est_bytes is None:
                 est_bytes = estimated
             else:
-                # Feed a real measurement back so the PRE-FIT admission
-                # gate can use it (safely, as a bound for smaller/equal
-                # future shapes -- see memory/adaptive.py). What admission
-                # must bound is the fit's transient PEAK (that's what
-                # OOMs), so prefer the peak hint when the backend reports
-                # one; the resident size is only a (low) fallback for
-                # backends without peak measurement.
+                # Feed real measurements back so the PRE-FIT admission gate
+                # can use them (safely, as a bound for smaller/equal future
+                # shapes -- see memory/adaptive.py). The estimator decides
+                # what the fit needed: resident + working floor for models
+                # whose working memory adapts to what is free (their peak
+                # only says how much was free), else the transient peak.
                 peak_hint = getattr(chosen, "fit_peak_bytes_hint", None)
                 peak_bytes = peak_hint() if callable(peak_hint) else None
-                estimator.record_observation(
-                    n_train, n_features, peak_bytes if peak_bytes else est_bytes
-                )
+                record_fit = getattr(estimator, "record_fit", None)
+                if callable(record_fit):
+                    record_fit(n_train, n_features, est_bytes, peak_bytes)
+                else:
+                    estimator.record_observation(
+                        n_train, n_features, peak_bytes if peak_bytes else est_bytes
+                    )
             context = CachedContext(
                 dataset_id=resolved_id,
                 backend_name=backend_name,
diff --git a/src/tabctx/memory/adaptive.py b/src/tabctx/memory/adaptive.py
index a776518..0049625 100644
--- a/src/tabctx/memory/adaptive.py
+++ b/src/tabctx/memory/adaptive.py
@@ -35,6 +35,7 @@
 from dataclasses import dataclass
 
 from tabctx.memory.estimator import MemoryEstimator
+from tabctx.memory.requirement import WorkingFloorModel
 
 # The n_test basis at which calibration predict-peaks were measured
 # (benchmarks/calibrate_memory.py's N_TEST_ROWS).
@@ -69,6 +70,7 @@ def __init__(
         # run-to-run peak variance eat several hundred MB, so the top ~15%
         # of the device is treated as unbudgetable.
         transient_capacity_fraction: float = 0.85,
+        working_floor: WorkingFloorModel | None = None,
     ) -> None:
         """preloaded: factory-installed observations (e.g. a measured
         calibration grid -- see memory/calibration_data.py), consulted
@@ -89,7 +91,12 @@ def __init__(
         transient_capacity_fraction: admission_headroom_bytes() allows a
         fit's estimated peak to use up to this fraction of the
         (fraction-scaled) device capacity MINUS what the cache already
-        holds resident -- peak-plus-resident is what actually OOMs."""
+        holds resident -- peak-plus-resident is what actually OOMs.
+
+        working_floor: when given (a model whose working memory adapts to
+        what is free, e.g. TabICL -- see memory/requirement.py), fits are
+        learned as resident + working floor, the memory they actually need,
+        instead of their unconstrained peak."""
         self._fallback = fallback
         self._safety_margin = safety_margin
         self._preloaded: tuple[Observation, ...] = tuple(preloaded)
@@ -102,6 +109,7 @@ def __init__(
         self._preloaded_predict: tuple[Observation, ...] = tuple(preloaded_predict)
         self._preloaded_margin = preloaded_margin
         self._transient_capacity_fraction = transient_capacity_fraction
+        self._working_floor = working_floor
         # Bounded FIFO -- simple, avoids unbounded growth over a long-
         # running replica's lifetime. A future version could prune to a
         # Pareto frontier of dominating points instead of dropping
@@ -112,9 +120,29 @@ def __init__(
     def record_observation(
         self, n_train: int, n_features: int, real_bytes: int
     ) -> None:
+        """Record a fit's measured requirement directly."""
         with self._lock:
             self._observations.append(Observation(n_train, n_features, real_bytes))
 
+    def record_fit(
+        self,
+        n_train: int,
+        n_features: int,
+        resident_bytes: int,
+        peak_bytes: int | None,
+    ) -> None:
+        """Learn from a completed fit. With a working-floor model the
+        requirement is resident + floor (the peak of an adaptive model only
+        says how much was free); without one it is the peak, falling back
+        to the resident size for backends that don't report peaks."""
+        if self._working_floor is not None:
+            need = resident_bytes + self._working_floor.bytes(n_train, n_features)
+            if peak_bytes:
+                need = min(need, peak_bytes)  # never more than it actually used
+        else:
+            need = peak_bytes or resident_bytes
+        self.record_observation(n_train, n_features, need)
+
     def _best_fit_estimate(self, n_train: int, n_features: int) -> int | None:
         """The MINIMUM effective estimate (measurement x its margin)
         across dominating measurements, or None.
diff --git a/src/tabctx/memory/calibration_tabicl_a100_requirement.py b/src/tabctx/memory/calibration_tabicl_a100_requirement.py
new file mode 100644
index 0000000..268d93f
--- /dev/null
+++ b/src/tabctx/memory/calibration_tabicl_a100_requirement.py
@@ -0,0 +1,81 @@
+"""GENERATED by benchmarks/gen_requirement_grid.py -- do not hand-edit.
+
+Measured requirements (NVIDIA A100-SXM4-40GB, tabicl 2.2.0, kv mode, native KV precision, 2026-10-17): for each shape, the least free GPU
+memory with which a fit succeeded (resident context + working floor), from
+benchmarks/bench_memory_pressure.py --search. See memory/requirement.py.
+"""
+
+from __future__ import annotations
+
+from tabctx.memory.adaptive import Observation
+from tabctx.memory.requirement import FloorPoint
+
+A100_40GB_TABICL_KV_FLOOR_POINTS: tuple[FloorPoint, ...] = (
+    FloorPoint(
+        n_train=2000, n_features=10, working_bytes=322122547
+    ),  # resident 0.39 GiB, unconstrained peak 0.59 GiB, fit 0.433s -> 0.277s at the floor
+    FloorPoint(
+        n_train=20000, n_features=10, working_bytes=1760936591
+    ),  # resident 3.68 GiB, unconstrained peak 5.53 GiB, fit 2.141s -> 1.327s at the floor
+    FloorPoint(
+        n_train=45000, n_features=10, working_bytes=8632884264
+    ),  # resident 6.58 GiB, unconstrained peak 10.72 GiB, fit 4.581s -> 4.051s at the floor
+    FloorPoint(
+        n_train=80000, n_features=10, working_bytes=10157597655
+    ),  # resident 11.23 GiB, unconstrained peak 18.59 GiB, fit 10.753s -> 9.878s at the floor
+    FloorPoint(
+        n_train=2000, n_features=50, working_bytes=644245094
+    ),  # resident 0.45 GiB, unconstrained peak 1.64 GiB, fit 0.694s -> 0.74s at the floor
+    FloorPoint(
+        n_train=10000, n_features=50, working_bytes=923417968
+    ),  # resident 1.92 GiB, unconstrained peak 7.76 GiB, fit 1.957s -> 1.764s at the floor
+    FloorPoint(
+        n_train=40000, n_features=50, working_bytes=4724464025
+    ),  # resident 7.4 GiB, unconstrained peak 25.85 GiB, fit 6.619s -> 8.401s at the floor
+    FloorPoint(
+        n_train=2000, n_features=100, working_bytes=644245094
+    ),  # resident 0.52 GiB, unconstrained peak 2.98 GiB, fit 1.273s -> 1.707s at the floor
+    FloorPoint(
+        n_train=20000, n_features=100, working_bytes=1760936591
+    ),  # resident 3.81 GiB, unconstrained peak 24.94 GiB, fit 5.871s -> 8.606s at the floor
+    FloorPoint(
+        n_train=50000, n_features=100, working_bytes=5712306503
+    ),  # resident 7.72 GiB, unconstrained peak 30.15 GiB, fit 13.99s -> 21.466s at the floor
+    FloorPoint(
+        n_train=1000, n_features=300, working_bytes=633507676
+    ),  # resident 0.63 GiB, unconstrained peak 4.48 GiB, fit 3.021s -> 3.665s at the floor
+    FloorPoint(
+        n_train=5000, n_features=300, working_bytes=934155386
+    ),  # resident 1.38 GiB, unconstrained peak 20.0 GiB, fit 5.614s -> 8.485s at the floor
+    FloorPoint(
+        n_train=1000, n_features=500, working_bytes=633507676
+    ),  # resident 0.92 GiB, unconstrained peak 7.36 GiB, fit 5.105s -> 6.22s at the floor
+    FloorPoint(
+        n_train=5000, n_features=500, working_bytes=1245540515
+    ),  # resident 1.68 GiB, unconstrained peak 25.27 GiB, fit 9.848s -> 13.435s at the floor
+    FloorPoint(
+        n_train=500, n_features=1000, working_bytes=934155386
+    ),  # resident 1.56 GiB, unconstrained peak 8.26 GiB, fit 9.29s -> 9.966s at the floor
+    FloorPoint(
+        n_train=2000, n_features=1000, working_bytes=1234803097
+    ),  # resident 1.84 GiB, unconstrained peak 24.56 GiB, fit 13.093s -> 16.456s at the floor
+)
+
+A100_40GB_TABICL_KV_REQUIREMENT_GRID: tuple[Observation, ...] = (
+    Observation(n_train=2000, n_features=10, real_bytes=740881858),
+    Observation(n_train=20000, n_features=10, real_bytes=5712306503),
+    Observation(n_train=45000, n_features=10, real_bytes=15698105466),
+    Observation(n_train=80000, n_features=10, real_bytes=22215718338),
+    Observation(n_train=2000, n_features=50, real_bytes=1127428915),
+    Observation(n_train=10000, n_features=50, real_bytes=2985002270),
+    Observation(n_train=40000, n_features=50, real_bytes=12670153523),
+    Observation(n_train=2000, n_features=100, real_bytes=1202590842),
+    Observation(n_train=20000, n_features=100, real_bytes=5851892940),
+    Observation(n_train=50000, n_features=100, real_bytes=14001593384),
+    Observation(n_train=1000, n_features=300, real_bytes=1309965025),
+    Observation(n_train=5000, n_features=300, real_bytes=2415919104),
+    Observation(n_train=1000, n_features=500, real_bytes=1621350154),
+    Observation(n_train=5000, n_features=500, real_bytes=3049426780),
+    Observation(n_train=500, n_features=1000, real_bytes=2609192632),
+    Observation(n_train=2000, n_features=1000, real_bytes=3210488053),
+)
diff --git a/src/tabctx/memory/requirement.py b/src/tabctx/memory/requirement.py
new file mode 100644
index 0000000..bb590e3
--- /dev/null
+++ b/src/tabctx/memory/requirement.py
@@ -0,0 +1,72 @@
+"""What a fit actually NEEDS on the device: resident context + working floor.
+
+Measured on an A100 (benchmarks/bench_memory_pressure.py --search;
+benchmarks/baselines/minfree-tabicl-kv-a100.jsonl): TabICL sizes its
+stage batches from whatever GPU memory is free, so an unconstrained fit's
+transient peak grows to fill the card. It does not reflect a requirement.
+With only a fixed amount left free, the same fits succeed using a small
+*working floor* on top of the resident context they build:
+
+    5,000 x 500    peak 25.3 GB   needs 1.7 GB resident + 1.2 GB working
+    20,000 x 100   peak 24.9 GB   needs 3.8 GB resident + 1.6 GB working
+    80,000 x 10    peak 18.6 GB   needs 11.2 GB resident + 9.5 GB working
+
+at 1-1.5x the fit time. So admission should bound resident + floor, not
+the peak: bounding the peak rejects tables that fit and evicts other
+tenants' contexts to make room the fit doesn't need. Within that room,
+the model still uses whatever extra memory is free.
+
+``WorkingFloorModel`` turns measured floors into a conservative estimate
+for any shape: the smallest margin-scaled floor among measurements that
+dominate the shape (at least as many rows and features -- the floor grows
+with both), else a per-row extrapolation fit as an upper envelope of the
+measurements.
+"""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Iterable
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class FloorPoint:
+    """A measured working floor: the least memory beyond the resident
+    context with which a fit of this shape succeeded."""
+
+    n_train: int
+    n_features: int
+    working_bytes: int
+
+    def dominates(self, n_train: int, n_features: int) -> bool:
+        return self.n_train >= n_train and self.n_features >= n_features
+
+
+class WorkingFloorModel:
+    def __init__(
+        self,
+        points: Iterable[FloorPoint],
+        margin: float = 1.25,
+        base_bytes: int = 1 << 30,
+    ) -> None:
+        """margin: applied to measured floors (the binary search has finite
+        precision, and fragmentation varies run to run). base_bytes: the
+        floor's intercept for the extrapolation."""
+        self._points = tuple(points)
+        if not self._points:
+            raise ValueError("a working floor model needs measurements")
+        self._margin = margin
+        self._base = base_bytes
+        # Upper envelope through the intercept: the steepest per-row slope
+        # any measurement needs, so extrapolation never undershoots one.
+        self._per_row = max(
+            max(0, p.working_bytes * margin - base_bytes) / p.n_train
+            for p in self._points
+        )
+
+    def bytes(self, n_train: int, n_features: int) -> int:
+        dominating = [p for p in self._points if p.dominates(n_train, n_features)]
+        if dominating:
+            return math.ceil(min(p.working_bytes for p in dominating) * self._margin)
+        return math.ceil(self._base + self._per_row * n_train)
diff --git a/src/tabctx/serve/factory.py b/src/tabctx/serve/factory.py
index b37994e..f3a6c3c 100644
--- a/src/tabctx/serve/factory.py
+++ b/src/tabctx/serve/factory.py
@@ -80,6 +80,7 @@
     DEFAULT_HARD_CEILING_BYTES,
     MemoryEstimator,
 )
+from tabctx.memory.requirement import WorkingFloorModel
 from tabctx.multidevice import MultiDeviceEngine
 from tabctx.serve.uploads import DEFAULT_MAX_UPLOAD_BYTES, DEFAULT_TTL_S
 
@@ -284,6 +285,26 @@ def _preloaded_observations(
     return fit_grid, predict_grid
 
 
+def _requirement_calibration(
+    kind: BackendKind, kv_cache: KvCacheMode
+) -> tuple[WorkingFloorModel | None, tuple]:
+    """Measured working floors + requirements for models whose working
+    memory adapts to what is free (memory/requirement.py), so admission
+    bounds resident + floor instead of an unconstrained peak. Measured for
+    TabICL in kv mode so far; other configurations keep peak semantics."""
+    if kind != "tabicl" or kv_cache != "kv":
+        return None, ()
+    try:
+        from tabctx.memory import calibration_tabicl_a100_requirement as grid
+    except ImportError:  # generated module absent: peak semantics, no preload
+        return None, ()
+
+    return (
+        WorkingFloorModel(grid.A100_40GB_TABICL_KV_FLOOR_POINTS),
+        grid.A100_40GB_TABICL_KV_REQUIREMENT_GRID,
+    )
+
+
 def detect_gpu_capacity_bytes(devices: tuple[str, ...] = ("cuda",)) -> int | None:
     """Total memory of the given CUDA devices (the smallest, since one
     estimator budgets them all), or None when torch is absent or sees no
@@ -368,10 +389,13 @@ def build_estimator(
         gpu_capacity_bytes=int(capacity * fraction),
     )
     fit_grid, predict_grid = _preloaded_observations(kind, settings.kv_cache)
+    working_floor, requirements = _requirement_calibration(kind, settings.kv_cache)
     return AdaptiveMemoryEstimator(
         fallback=fallback,
-        preloaded=fit_grid,
+        # Both are valid upper bounds; the tightest dominating one wins.
+        preloaded=tuple(fit_grid) + tuple(requirements),
         preloaded_predict=predict_grid,
+        working_floor=working_floor,
     )
 
 
diff --git a/tests/unit/test_requirement.py b/tests/unit/test_requirement.py
new file mode 100644
index 0000000..2b139a5
--- /dev/null
+++ b/tests/unit/test_requirement.py
@@ -0,0 +1,122 @@
+"""Unit tests for requirement-based admission (memory/requirement.py):
+fits are learned as resident + working floor, not unconstrained peaks."""
+
+import pytest
+
+from tabctx.backends.fake import FakeBackend
+from tabctx.cache.manager import CachedContext, ContextCacheManager
+from tabctx.engine import TabctxEngine
+from tabctx.memory import (
+    A100_40GB_TABICL_CALIBRATION,
+    AdaptiveMemoryEstimator,
+    PowerLawMemoryEstimator,
+)
+from tabctx.memory.adaptive import Observation
+from tabctx.memory.requirement import FloorPoint, WorkingFloorModel
+
+GiB = 1024**3
+POINTS = (
+    FloorPoint(2_000, 10, int(0.3 * GiB)),
+    FloorPoint(20_000, 100, int(1.6 * GiB)),
+    FloorPoint(80_000, 10, int(9.5 * GiB)),
+)
+
+
+class TestWorkingFloorModel:
+    def test_needs_measurements(self):
+        with pytest.raises(ValueError):
+            WorkingFloorModel(())
+
+    def test_dominated_shape_uses_the_tightest_measurement(self):
+        m = WorkingFloorModel(POINTS, margin=1.25)
+        # dominated by (20000,100) and (80000,10)? only (20000,100) has >=50 features
+        assert m.bytes(10_000, 50) == pytest.approx(1.6 * GiB * 1.25, rel=1e-6)
+        assert m.bytes(1_000, 5) == pytest.approx(0.3 * GiB * 1.25, rel=1e-6)
+
+    def test_extrapolation_never_undershoots_a_measurement(self):
+        m = WorkingFloorModel(POINTS, margin=1.25)
+        for p in POINTS:
+            # a slightly larger shape than any measurement: extrapolated
+            assert m.bytes(p.n_train + 1, 10_000) >= p.working_bytes * 1.25
+
+    def test_extrapolation_grows_with_rows(self):
+        m = WorkingFloorModel(POINTS)
+        assert m.bytes(200_000, 10) > m.bytes(100_000, 10) > m.bytes(90_000, 10)
+
+
+def _estimator(floor=None, capacity=40 * GiB, preloaded=()):
+    fallback = PowerLawMemoryEstimator(
+        A100_40GB_TABICL_CALIBRATION,
+        hard_ceiling_bytes=int(capacity * 0.6),
+        gpu_capacity_bytes=capacity,
+    )
+    return AdaptiveMemoryEstimator(
+        fallback=fallback, working_floor=floor, preloaded=preloaded
+    )
+
+
+class TestRecordFit:
+    def test_with_a_floor_model_learns_resident_plus_floor(self):
+        floor = WorkingFloorModel(POINTS, margin=1.0)
+        est = _estimator(floor)
+        est.record_fit(20_000, 100, resident_bytes=4 * GiB, peak_bytes=25 * GiB)
+        expected = (4 + 1.6) * GiB * 1.5  # runtime safety margin
+        assert est.estimate_bytes(20_000, 0, 100) == pytest.approx(expected, rel=1e-3)
+
+    def test_never_more_than_the_measured_peak(self):
+        est = _estimator(WorkingFloorModel(POINTS, margin=1.0))
+        est.record_fit(20_000, 100, resident_bytes=4 * GiB, peak_bytes=int(4.5 * GiB))
+        assert est.estimate_bytes(20_000, 0, 100) == pytest.approx(
+            4.5 * GiB * 1.5, rel=1e-3
+        )
+
+    def test_without_a_floor_model_learns_the_peak(self):
+        est = _estimator()
+        est.record_fit(20_000, 100, resident_bytes=4 * GiB, peak_bytes=25 * GiB)
+        assert est.estimate_bytes(20_000, 0, 100) == pytest.approx(
+            25 * GiB * 1.5, rel=1e-3
+        )
+        est2 = _estimator()
+        est2.record_fit(1_000, 10, resident_bytes=GiB, peak_bytes=None)
+        assert est2.estimate_bytes(1_000, 0, 10) == pytest.approx(GiB * 1.5, rel=1e-3)
+
+
+class _AdaptiveFake(FakeBackend):
+    """Reports a resident context and an unconstrained peak that fills the
+    card, like TabICL with plenty free."""
+
+    def __init__(self, resident: int, peak: int) -> None:
+        super().__init__(bytes_hint=resident, peak_bytes_hint=peak)
+
+
+class TestAdmissionEffect:
+    def test_wide_table_is_admitted_without_evicting_other_tenants(self):
+        """The motivating case: a 20k x 100 fit's unconstrained peak is 25 GB,
+        but it needs only resident + floor (5.6 GB). Each estimator is
+        calibrated as the factory does (a measured requirement vs a measured
+        peak). Beside a warm cache of five 4 GB tenants, peak semantics evict
+        tenants to make room the fit doesn't need; requirement semantics
+        admit it untouched."""
+
+        def run(floor):
+            measured = (5.6 if floor else 25) * GiB
+            est = _estimator(
+                floor, preloaded=(Observation(20_000, 100, int(measured)),)
+            )
+            backend = _AdaptiveFake(resident=4 * GiB, peak=25 * GiB)
+            cache = ContextCacheManager(capacity_bytes=34 * GiB)
+            engine = TabctxEngine(backend=backend, cache=cache, estimator=est)
+            X = [[0.0] * 100] * 20_000
+            y = ["a"] * 20_000
+            engine.fit(X, y, dataset_id="first")  # learn from one fit
+            for i in range(5):  # other tenants: 5 x 4 GB resident
+                cache.put(
+                    CachedContext(
+                        f"t{i}", "fake", "classification", 10, 1, None, 4 * GiB
+                    )
+                )
+            engine.fit(X, y, dataset_id="second")
+            return [f"t{i}" for i in range(5) if cache.holds(f"t{i}")]
+
+        assert len(run(WorkingFloorModel(POINTS))) == 5  # nobody evicted
+        assert len(run(None)) < 5  # peak semantics evict tenants
diff --git a/tests/unit/test_serve_factory.py b/tests/unit/test_serve_factory.py
index 4a07467..0cc0939 100644
--- a/tests/unit/test_serve_factory.py
+++ b/tests/unit/test_serve_factory.py
@@ -319,8 +319,9 @@ def test_missing_calibration_module_degrades_to_no_preload(self, monkeypatch):
         # shortcut finds the real module other tests already imported.
         import tabctx.memory as memory_pkg
 
-        monkeypatch.delattr(memory_pkg, "calibration_tabicl_a100", raising=False)
-        monkeypatch.setitem(sys.modules, "tabctx.memory.calibration_tabicl_a100", None)
+        for name in ("calibration_tabicl_a100", "calibration_tabicl_a100_requirement"):
+            monkeypatch.delattr(memory_pkg, name, raising=False)
+            monkeypatch.setitem(sys.modules, f"tabctx.memory.{name}", None)
         est = build_estimator(ServeSettings(backends=("tabicl",)))
         assert "0 preloaded calibration measurement" in est.confidence()
 

From 8167b5ad2f8777ad37925580182334d3921e7811 Mon Sep 17 00:00:00 2001
From: Amrit Krishnan <amrit110@gmail.com>
Date: Sat, 17 Oct 2026 16:02:00 -0400
Subject: [PATCH 04/39] feat(memory): extrapolate uncalibrated shapes as
 resident (linear in rows) + floor

ResidentRateModel: a kv context is a feature-dependent column cache plus
one ICL key/value per row, so resident(n,d) <= r0 + slope*(n-n0) from any
measured fit with >= d features; slope from long fits (runtime fits
supersede calibration). Replaces the fallback formula for shapes beyond
the grid (130k x 100: 165 GiB -> 43 GiB). Floors are bounded by what the
fit used unconstrained (a ballast-fragmentation artifact inflated one);
floor slope from long fits only; wider-than-measured shapes fall back.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 benchmarks/gen_requirement_grid.py            | 11 ++-
 benchmarks/sweep_engine.py                    | 76 ++++++++++++++++++
 src/tabctx/memory/adaptive.py                 | 36 ++++++++-
 .../calibration_tabicl_a100_requirement.py    | 72 +++++++++++------
 src/tabctx/memory/requirement.py              | 80 +++++++++++++++++--
 src/tabctx/serve/factory.py                   | 20 ++---
 tests/unit/test_requirement.py                | 10 ++-
 7 files changed, 255 insertions(+), 50 deletions(-)
 create mode 100644 benchmarks/sweep_engine.py

diff --git a/benchmarks/gen_requirement_grid.py b/benchmarks/gen_requirement_grid.py
index b986c15..26a3637 100644
--- a/benchmarks/gen_requirement_grid.py
+++ b/benchmarks/gen_requirement_grid.py
@@ -40,15 +40,20 @@ def main() -> int:
     floors, needs = [], []
     for r in ok:
         n, d = r["rows"], r["features"]
-        working = int(r["min_working_gib"] * GiB)
-        need = int(r["min_free_gib"] * GiB)
+        resident = int(r["resident_gib"] * GiB)
+        # The floor can't exceed what the fit used unconstrained: a search
+        # result above that is ballast fragmentation, not need.
+        used = int((r["unconstrained_peak_gib"] - r["resident_gib"]) * GiB)
+        working = min(int(r["min_working_gib"] * GiB), used)
+        need = resident + working
         note = (
             f"resident {r['resident_gib']} GiB, unconstrained peak "
             f"{r['unconstrained_peak_gib']} GiB, fit {r['unconstrained_fit_s']}s "
             f"-> {r['fit_s_at_min']}s at the floor"
         )
         floors.append(
-            f"    FloorPoint(n_train={n}, n_features={d}, working_bytes={working}),  # {note}"
+            f"    FloorPoint(n_train={n}, n_features={d}, working_bytes={working}, "
+            f"resident_bytes={resident}),  # {note}"
         )
         needs.append(
             f"    Observation(n_train={n}, n_features={d}, real_bytes={need}),"
diff --git a/benchmarks/sweep_engine.py b/benchmarks/sweep_engine.py
new file mode 100644
index 0000000..81c01f5
--- /dev/null
+++ b/benchmarks/sweep_engine.py
@@ -0,0 +1,76 @@
+#!/usr/bin/env python3
+"""Admission sweep through the engine on one GPU, no serving layer: for
+each feature count, fit tables of increasing rows (re-fitting one id) until
+the admission gate rejects, recording fit time, peak, resident and any
+OOM. Every admitted fit must succeed -- an OOM means admission is too
+optimistic. Compare admission policies with --peak-only (the pre-floor
+behavior: no working-floor model, peak grids only).
+
+Usage: sweep_engine.py --backend tabicl --features 10 100 500 1000 --out sweep.jsonl
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+import time
+
+import numpy as np
+
+import tabctx.serve.factory as factory
+from tabctx.errors import AdmissionRejected, BackendComputeError
+from tabctx.serve.factory import ServeSettings, build_engine
+
+ROWS = [1_000, 2_000, 5_000, 10_000, 20_000, 30_000, 45_000, 65_000, 80_000,
+        100_000, 130_000, 160_000, 200_000, 250_000]  # fmt: skip
+
+
+def main() -> int:
+    p = argparse.ArgumentParser()
+    p.add_argument("--backend", default="tabicl")
+    p.add_argument("--features", type=int, nargs="+", default=[10, 100, 500, 1000])
+    p.add_argument("--peak-only", action="store_true")
+    p.add_argument("--out", default=None)
+    args = p.parse_args()
+    if args.peak_only:
+        factory._requirement_calibration = lambda kind, kv: (None, (), ())
+    built = build_engine(ServeSettings(backends=(args.backend,)))
+    engine = built.engine
+    rng = np.random.default_rng(0)
+    out = open(args.out, "a") if args.out else None  # noqa: SIM115
+    policy = "peak-only" if args.peak_only else "requirement"
+    for d in args.features:
+        for n in ROWS:
+            X = rng.normal(size=(n, d)).astype(np.float32)
+            y = (X[:, 0] > 0).astype(int)
+            rec = {"policy": policy, "rows": n, "features": d}
+            t0 = time.perf_counter()
+            try:
+                engine.fit(X, y, dataset_id=f"sweep-{d}")
+                rec["status"] = "ok"
+            except AdmissionRejected as e:
+                rec["status"] = "413"
+                rec["detail"] = str(e)[:160]
+            except BackendComputeError as e:
+                rec["status"] = "OOM"
+                rec["detail"] = str(e)[:160]
+            rec["fit_s"] = round(time.perf_counter() - t0, 2)
+            backend = built.backend
+            if rec["status"] == "ok":
+                rec["peak_gib"] = round((backend.fit_peak_bytes_hint() or 0) / 2**30, 2)
+                rec["resident_gib"] = round(
+                    (backend.context_bytes_hint(n, d) or 0) / 2**30, 2
+                )
+            print(json.dumps(rec), flush=True)
+            if out:
+                out.write(json.dumps(rec) + "\n")
+                out.flush()
+            if rec["status"] != "ok":
+                break
+        engine.evict(f"sweep-{d}")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/src/tabctx/memory/adaptive.py b/src/tabctx/memory/adaptive.py
index 0049625..30adf3e 100644
--- a/src/tabctx/memory/adaptive.py
+++ b/src/tabctx/memory/adaptive.py
@@ -35,7 +35,7 @@
 from dataclasses import dataclass
 
 from tabctx.memory.estimator import MemoryEstimator
-from tabctx.memory.requirement import WorkingFloorModel
+from tabctx.memory.requirement import ResidentRateModel, WorkingFloorModel
 
 # The n_test basis at which calibration predict-peaks were measured
 # (benchmarks/calibrate_memory.py's N_TEST_ROWS).
@@ -71,6 +71,8 @@ def __init__(
         # of the device is treated as unbudgetable.
         transient_capacity_fraction: float = 0.85,
         working_floor: WorkingFloorModel | None = None,
+        resident_rate: ResidentRateModel | None = None,
+        resident_margin: float = 1.2,
     ) -> None:
         """preloaded: factory-installed observations (e.g. a measured
         calibration grid -- see memory/calibration_data.py), consulted
@@ -96,7 +98,12 @@ def __init__(
         working_floor: when given (a model whose working memory adapts to
         what is free, e.g. TabICL -- see memory/requirement.py), fits are
         learned as resident + working floor, the memory they actually need,
-        instead of their unconstrained peak."""
+        instead of their unconstrained peak.
+
+        resident_rate: with working_floor, extrapolates shapes no
+        measurement dominates as (extrapolated resident x resident_margin)
+        + floor -- a kv context grows linearly in rows -- instead of the
+        fallback formula. Learns from every fit."""
         self._fallback = fallback
         self._safety_margin = safety_margin
         self._preloaded: tuple[Observation, ...] = tuple(preloaded)
@@ -110,6 +117,8 @@ def __init__(
         self._preloaded_margin = preloaded_margin
         self._transient_capacity_fraction = transient_capacity_fraction
         self._working_floor = working_floor
+        self._resident_rate = resident_rate
+        self._resident_margin = resident_margin
         # Bounded FIFO -- simple, avoids unbounded growth over a long-
         # running replica's lifetime. A future version could prune to a
         # Pareto frontier of dominating points instead of dropping
@@ -135,14 +144,30 @@ def record_fit(
         requirement is resident + floor (the peak of an adaptive model only
         says how much was free); without one it is the peak, falling back
         to the resident size for backends that don't report peaks."""
-        if self._working_floor is not None:
-            need = resident_bytes + self._working_floor.bytes(n_train, n_features)
+        if self._resident_rate is not None:
+            self._resident_rate.record(n_train, n_features, resident_bytes)
+        floor = (
+            self._working_floor.bytes(n_train, n_features)
+            if self._working_floor is not None
+            else None
+        )
+        if floor is not None:
+            need = resident_bytes + floor
             if peak_bytes:
                 need = min(need, peak_bytes)  # never more than it actually used
         else:
             need = peak_bytes or resident_bytes
         self.record_observation(n_train, n_features, need)
 
+    def _extrapolated_requirement(self, n_train: int, n_features: int) -> int | None:
+        if self._working_floor is None or self._resident_rate is None:
+            return None
+        resident = self._resident_rate.estimate(n_train, n_features)
+        floor = self._working_floor.bytes(n_train, n_features)
+        if resident is None or floor is None:
+            return None
+        return math.ceil(resident * self._resident_margin) + floor
+
     def _best_fit_estimate(self, n_train: int, n_features: int) -> int | None:
         """The MINIMUM effective estimate (measurement x its margin)
         across dominating measurements, or None.
@@ -169,6 +194,9 @@ def estimate_bytes(self, n_train: int, n_test: int, n_features: int) -> int:
             best = self._best_fit_estimate(n_train, n_features)
             if best is not None:
                 return best
+            extrapolated = self._extrapolated_requirement(n_train, n_features)
+            if extrapolated is not None:
+                return extrapolated
         else:
             # Predict-time (chunking) query: use a measured predict peak
             # when a dominating one exists, scaled linearly in n_test
diff --git a/src/tabctx/memory/calibration_tabicl_a100_requirement.py b/src/tabctx/memory/calibration_tabicl_a100_requirement.py
index 268d93f..5426b00 100644
--- a/src/tabctx/memory/calibration_tabicl_a100_requirement.py
+++ b/src/tabctx/memory/calibration_tabicl_a100_requirement.py
@@ -12,70 +12,94 @@
 
 A100_40GB_TABICL_KV_FLOOR_POINTS: tuple[FloorPoint, ...] = (
     FloorPoint(
-        n_train=2000, n_features=10, working_bytes=322122547
+        n_train=2000, n_features=10, working_bytes=214748364, resident_bytes=418759311
     ),  # resident 0.39 GiB, unconstrained peak 0.59 GiB, fit 0.433s -> 0.277s at the floor
     FloorPoint(
-        n_train=20000, n_features=10, working_bytes=1760936591
+        n_train=20000,
+        n_features=10,
+        working_bytes=1760936591,
+        resident_bytes=3951369912,
     ),  # resident 3.68 GiB, unconstrained peak 5.53 GiB, fit 2.141s -> 1.327s at the floor
     FloorPoint(
-        n_train=45000, n_features=10, working_bytes=8632884264
+        n_train=45000,
+        n_features=10,
+        working_bytes=4445291151,
+        resident_bytes=7065221201,
     ),  # resident 6.58 GiB, unconstrained peak 10.72 GiB, fit 4.581s -> 4.051s at the floor
     FloorPoint(
-        n_train=80000, n_features=10, working_bytes=10157597655
+        n_train=80000,
+        n_features=10,
+        working_bytes=7902739824,
+        resident_bytes=12058120683,
     ),  # resident 11.23 GiB, unconstrained peak 18.59 GiB, fit 10.753s -> 9.878s at the floor
     FloorPoint(
-        n_train=2000, n_features=50, working_bytes=644245094
+        n_train=2000, n_features=50, working_bytes=644245094, resident_bytes=483183820
     ),  # resident 0.45 GiB, unconstrained peak 1.64 GiB, fit 0.694s -> 0.74s at the floor
     FloorPoint(
-        n_train=10000, n_features=50, working_bytes=923417968
+        n_train=10000, n_features=50, working_bytes=923417968, resident_bytes=2061584302
     ),  # resident 1.92 GiB, unconstrained peak 7.76 GiB, fit 1.957s -> 1.764s at the floor
     FloorPoint(
-        n_train=40000, n_features=50, working_bytes=4724464025
+        n_train=40000,
+        n_features=50,
+        working_bytes=4724464025,
+        resident_bytes=7945689497,
     ),  # resident 7.4 GiB, unconstrained peak 25.85 GiB, fit 6.619s -> 8.401s at the floor
     FloorPoint(
-        n_train=2000, n_features=100, working_bytes=644245094
+        n_train=2000, n_features=100, working_bytes=644245094, resident_bytes=558345748
     ),  # resident 0.52 GiB, unconstrained peak 2.98 GiB, fit 1.273s -> 1.707s at the floor
     FloorPoint(
-        n_train=20000, n_features=100, working_bytes=1760936591
+        n_train=20000,
+        n_features=100,
+        working_bytes=1760936591,
+        resident_bytes=4090956349,
     ),  # resident 3.81 GiB, unconstrained peak 24.94 GiB, fit 5.871s -> 8.606s at the floor
     FloorPoint(
-        n_train=50000, n_features=100, working_bytes=5712306503
+        n_train=50000,
+        n_features=100,
+        working_bytes=5712306503,
+        resident_bytes=8289286881,
     ),  # resident 7.72 GiB, unconstrained peak 30.15 GiB, fit 13.99s -> 21.466s at the floor
     FloorPoint(
-        n_train=1000, n_features=300, working_bytes=633507676
+        n_train=1000, n_features=300, working_bytes=633507676, resident_bytes=676457349
     ),  # resident 0.63 GiB, unconstrained peak 4.48 GiB, fit 3.021s -> 3.665s at the floor
     FloorPoint(
-        n_train=5000, n_features=300, working_bytes=934155386
+        n_train=5000, n_features=300, working_bytes=934155386, resident_bytes=1481763717
     ),  # resident 1.38 GiB, unconstrained peak 20.0 GiB, fit 5.614s -> 8.485s at the floor
     FloorPoint(
-        n_train=1000, n_features=500, working_bytes=633507676
+        n_train=1000, n_features=500, working_bytes=633507676, resident_bytes=987842478
     ),  # resident 0.92 GiB, unconstrained peak 7.36 GiB, fit 5.105s -> 6.22s at the floor
     FloorPoint(
-        n_train=5000, n_features=500, working_bytes=1245540515
+        n_train=5000,
+        n_features=500,
+        working_bytes=1245540515,
+        resident_bytes=1803886264,
     ),  # resident 1.68 GiB, unconstrained peak 25.27 GiB, fit 9.848s -> 13.435s at the floor
     FloorPoint(
-        n_train=500, n_features=1000, working_bytes=934155386
+        n_train=500, n_features=1000, working_bytes=934155386, resident_bytes=1675037245
     ),  # resident 1.56 GiB, unconstrained peak 8.26 GiB, fit 9.29s -> 9.966s at the floor
     FloorPoint(
-        n_train=2000, n_features=1000, working_bytes=1234803097
+        n_train=2000,
+        n_features=1000,
+        working_bytes=1234803097,
+        resident_bytes=1975684956,
     ),  # resident 1.84 GiB, unconstrained peak 24.56 GiB, fit 13.093s -> 16.456s at the floor
 )
 
 A100_40GB_TABICL_KV_REQUIREMENT_GRID: tuple[Observation, ...] = (
-    Observation(n_train=2000, n_features=10, real_bytes=740881858),
+    Observation(n_train=2000, n_features=10, real_bytes=633507675),
     Observation(n_train=20000, n_features=10, real_bytes=5712306503),
-    Observation(n_train=45000, n_features=10, real_bytes=15698105466),
-    Observation(n_train=80000, n_features=10, real_bytes=22215718338),
-    Observation(n_train=2000, n_features=50, real_bytes=1127428915),
+    Observation(n_train=45000, n_features=10, real_bytes=11510512352),
+    Observation(n_train=80000, n_features=10, real_bytes=19960860507),
+    Observation(n_train=2000, n_features=50, real_bytes=1127428914),
     Observation(n_train=10000, n_features=50, real_bytes=2985002270),
-    Observation(n_train=40000, n_features=50, real_bytes=12670153523),
+    Observation(n_train=40000, n_features=50, real_bytes=12670153522),
     Observation(n_train=2000, n_features=100, real_bytes=1202590842),
     Observation(n_train=20000, n_features=100, real_bytes=5851892940),
     Observation(n_train=50000, n_features=100, real_bytes=14001593384),
     Observation(n_train=1000, n_features=300, real_bytes=1309965025),
-    Observation(n_train=5000, n_features=300, real_bytes=2415919104),
+    Observation(n_train=5000, n_features=300, real_bytes=2415919103),
     Observation(n_train=1000, n_features=500, real_bytes=1621350154),
-    Observation(n_train=5000, n_features=500, real_bytes=3049426780),
-    Observation(n_train=500, n_features=1000, real_bytes=2609192632),
+    Observation(n_train=5000, n_features=500, real_bytes=3049426779),
+    Observation(n_train=500, n_features=1000, real_bytes=2609192631),
     Observation(n_train=2000, n_features=1000, real_bytes=3210488053),
 )
diff --git a/src/tabctx/memory/requirement.py b/src/tabctx/memory/requirement.py
index bb590e3..416ebf2 100644
--- a/src/tabctx/memory/requirement.py
+++ b/src/tabctx/memory/requirement.py
@@ -21,11 +21,19 @@
 dominate the shape (at least as many rows and features -- the floor grows
 with both), else a per-row extrapolation fit as an upper envelope of the
 measurements.
+
+``ResidentRateModel`` extrapolates the resident context to shapes no
+measurement dominates. TabICL's resident context is a kv cache with one
+key/value per training row (plus a column cache independent of rows), so
+it is linear in rows: resident(n, d) <= n * max rate observed at >= d
+features. Rates are per-row bytes from real fits, which take precedence
+over calibration points (calibration may be at another precision).
 """
 
 from __future__ import annotations
 
 import math
+import threading
 from collections.abc import Iterable
 from dataclasses import dataclass
 
@@ -33,11 +41,13 @@
 @dataclass(frozen=True)
 class FloorPoint:
     """A measured working floor: the least memory beyond the resident
-    context with which a fit of this shape succeeded."""
+    context with which a fit of this shape succeeded (and that resident
+    context's size, when recorded)."""
 
     n_train: int
     n_features: int
     working_bytes: int
+    resident_bytes: int | None = None
 
     def dominates(self, n_train: int, n_features: int) -> bool:
         return self.n_train >= n_train and self.n_features >= n_features
@@ -49,24 +59,82 @@ def __init__(
         points: Iterable[FloorPoint],
         margin: float = 1.25,
         base_bytes: int = 1 << 30,
+        slope_min_rows: int = 10_000,
     ) -> None:
         """margin: applied to measured floors (the binary search has finite
         precision, and fragmentation varies run to run). base_bytes: the
-        floor's intercept for the extrapolation."""
+        floor's intercept for row extrapolation. slope_min_rows: only fits
+        at least this long set the per-row slope -- shorter ones are
+        dominated by the intercept."""
         self._points = tuple(points)
         if not self._points:
             raise ValueError("a working floor model needs measurements")
         self._margin = margin
         self._base = base_bytes
+        self._max_features = max(p.n_features for p in self._points)
+        long_fits = [p for p in self._points if p.n_train >= slope_min_rows]
         # Upper envelope through the intercept: the steepest per-row slope
-        # any measurement needs, so extrapolation never undershoots one.
+        # a long fit needs, so row extrapolation never undershoots one.
         self._per_row = max(
-            max(0, p.working_bytes * margin - base_bytes) / p.n_train
-            for p in self._points
+            max(0.0, p.working_bytes * margin - base_bytes) / p.n_train
+            for p in (long_fits or self._points)
         )
 
-    def bytes(self, n_train: int, n_features: int) -> int:
+    def bytes(self, n_train: int, n_features: int) -> int | None:
+        """Conservative working floor, or None for shapes wider than any
+        measurement (the floor's growth with features isn't extrapolated)."""
         dominating = [p for p in self._points if p.dominates(n_train, n_features)]
         if dominating:
             return math.ceil(min(p.working_bytes for p in dominating) * self._margin)
+        if n_features > self._max_features:
+            return None
         return math.ceil(self._base + self._per_row * n_train)
+
+
+class ResidentRateModel:
+    """Extrapolates a kv context's resident size to longer tables.
+
+    A TabICL kv context is a column cache that depends on features but not
+    rows, plus an ICL cache with one key/value per training row whose width
+    doesn't depend on features. So for any measured fit (n0 rows, d0 >= d
+    features, resident r0), resident(n, d) <= r0 + rate * (n - n0) for
+    n >= n0, where rate is the steepest per-row slope among long fits
+    (>= slope_min_rows; short fits are dominated by the column cache).
+    The estimate is the smallest such bound. This replica's own fits feed
+    both the anchors and the slope, and supersede calibration slopes once
+    any long fit exists: calibration may be at another KV precision."""
+
+    def __init__(
+        self, calibration: Iterable[FloorPoint] = (), slope_min_rows: int = 10_000
+    ) -> None:
+        self._min_rows = slope_min_rows
+        self._calibration = [
+            (p.n_train, p.n_features, p.resident_bytes)
+            for p in calibration
+            if p.resident_bytes
+        ]
+        self._runtime: list[tuple[int, int, int]] = []
+        self._lock = threading.Lock()
+
+    def record(self, n_train: int, n_features: int, resident_bytes: int) -> None:
+        if n_train > 0 and resident_bytes > 0:
+            with self._lock:
+                self._runtime.append((n_train, n_features, resident_bytes))
+                del self._runtime[:-500]  # bounded, newest kept
+
+    def _slope(self, points: list[tuple[int, int, int]]) -> float | None:
+        long_fits = [res / n for n, _, res in points if n >= self._min_rows]
+        return max(long_fits) if long_fits else None
+
+    def estimate(self, n_train: int, n_features: int) -> int | None:
+        with self._lock:
+            runtime = list(self._runtime)
+        slope = self._slope(runtime) or self._slope(self._calibration)
+        if slope is None:
+            return None
+        bounds = [
+            res + slope * max(0, n_train - n)
+            for n, d, res in runtime + self._calibration
+            if d >= n_features
+        ]
+        return math.ceil(min(bounds)) if bounds else None
diff --git a/src/tabctx/serve/factory.py b/src/tabctx/serve/factory.py
index f3a6c3c..a40cd9a 100644
--- a/src/tabctx/serve/factory.py
+++ b/src/tabctx/serve/factory.py
@@ -80,7 +80,7 @@
     DEFAULT_HARD_CEILING_BYTES,
     MemoryEstimator,
 )
-from tabctx.memory.requirement import WorkingFloorModel
+from tabctx.memory.requirement import ResidentRateModel, WorkingFloorModel
 from tabctx.multidevice import MultiDeviceEngine
 from tabctx.serve.uploads import DEFAULT_MAX_UPLOAD_BYTES, DEFAULT_TTL_S
 
@@ -287,22 +287,19 @@ def _preloaded_observations(
 
 def _requirement_calibration(
     kind: BackendKind, kv_cache: KvCacheMode
-) -> tuple[WorkingFloorModel | None, tuple]:
+) -> tuple[WorkingFloorModel | None, tuple, tuple]:
     """Measured working floors + requirements for models whose working
     memory adapts to what is free (memory/requirement.py), so admission
     bounds resident + floor instead of an unconstrained peak. Measured for
     TabICL in kv mode so far; other configurations keep peak semantics."""
     if kind != "tabicl" or kv_cache != "kv":
-        return None, ()
+        return None, (), ()
     try:
         from tabctx.memory import calibration_tabicl_a100_requirement as grid
     except ImportError:  # generated module absent: peak semantics, no preload
-        return None, ()
-
-    return (
-        WorkingFloorModel(grid.A100_40GB_TABICL_KV_FLOOR_POINTS),
-        grid.A100_40GB_TABICL_KV_REQUIREMENT_GRID,
-    )
+        return None, (), ()
+    points = grid.A100_40GB_TABICL_KV_FLOOR_POINTS
+    return WorkingFloorModel(points), grid.A100_40GB_TABICL_KV_REQUIREMENT_GRID, points
 
 
 def detect_gpu_capacity_bytes(devices: tuple[str, ...] = ("cuda",)) -> int | None:
@@ -389,13 +386,16 @@ def build_estimator(
         gpu_capacity_bytes=int(capacity * fraction),
     )
     fit_grid, predict_grid = _preloaded_observations(kind, settings.kv_cache)
-    working_floor, requirements = _requirement_calibration(kind, settings.kv_cache)
+    working_floor, requirements, floor_points = _requirement_calibration(
+        kind, settings.kv_cache
+    )
     return AdaptiveMemoryEstimator(
         fallback=fallback,
         # Both are valid upper bounds; the tightest dominating one wins.
         preloaded=tuple(fit_grid) + tuple(requirements),
         preloaded_predict=predict_grid,
         working_floor=working_floor,
+        resident_rate=ResidentRateModel(floor_points) if working_floor else None,
     )
 
 
diff --git a/tests/unit/test_requirement.py b/tests/unit/test_requirement.py
index 2b139a5..debb5a0 100644
--- a/tests/unit/test_requirement.py
+++ b/tests/unit/test_requirement.py
@@ -33,11 +33,15 @@ def test_dominated_shape_uses_the_tightest_measurement(self):
         assert m.bytes(10_000, 50) == pytest.approx(1.6 * GiB * 1.25, rel=1e-6)
         assert m.bytes(1_000, 5) == pytest.approx(0.3 * GiB * 1.25, rel=1e-6)
 
-    def test_extrapolation_never_undershoots_a_measurement(self):
+    def test_row_extrapolation_never_undershoots_a_long_fit(self):
         m = WorkingFloorModel(POINTS, margin=1.25)
         for p in POINTS:
-            # a slightly larger shape than any measurement: extrapolated
-            assert m.bytes(p.n_train + 1, 10_000) >= p.working_bytes * 1.25
+            if p.n_train >= 10_000:  # long fits set the slope
+                # more rows than any measurement at these features
+                assert m.bytes(10 * p.n_train, p.n_features) >= p.working_bytes * 1.25
+
+    def test_wider_than_any_measurement_is_unknown(self):
+        assert WorkingFloorModel(POINTS).bytes(100, 5_000) is None
 
     def test_extrapolation_grows_with_rows(self):
         m = WorkingFloorModel(POINTS)

From fa6bdca71f82b8820509d1c0cd0046a96df372ce Mon Sep 17 00:00:00 2001
From: Amrit Krishnan <amrit110@gmail.com>
Date: Sat, 17 Oct 2026 16:31:00 -0400
Subject: [PATCH 05/39] fix(tabicl): release pinned host memory after calls;
 host-memory admission guard

Large fits offload stage outputs to pinned host memory. The shared
backbone kept its PinnedBufferPool and the CUDA host allocator cached
the rest: +43 GiB host RSS retained after one 20000x500 fit (A100), so
two large tables OOM-killed a 64 GiB pod. Pools are cleared and the host
cache emptied after every fit/predict: retained 43 GiB -> 0.3 GiB.
Transient peak is still ~46 GiB per large fit, so HostMemoryGuard
(memory/host.py, process-wide, shared by every device) caps concurrent
large fits at host budget // 48 GiB; small fits never wait, and a
waiting fit does not block its device. /readyz reports slots.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 benchmarks/baselines/host-tabicl-a100.jsonl |  12 ++
 benchmarks/bench_host_memory.py             | 128 ++++++++++++++++++++
 benchmarks/gke/devbox.sh                    |   9 +-
 src/tabctx/backends/tabicl.py               |  27 +++++
 src/tabctx/engine.py                        |  15 ++-
 src/tabctx/memory/host.py                   |  95 +++++++++++++++
 src/tabctx/serve/app.py                     |   2 +
 src/tabctx/serve/factory.py                 |   9 ++
 tests/unit/test_host_memory.py              |  99 +++++++++++++++
 9 files changed, 394 insertions(+), 2 deletions(-)
 create mode 100644 benchmarks/baselines/host-tabicl-a100.jsonl
 create mode 100644 benchmarks/bench_host_memory.py
 create mode 100644 src/tabctx/memory/host.py
 create mode 100644 tests/unit/test_host_memory.py

diff --git a/benchmarks/baselines/host-tabicl-a100.jsonl b/benchmarks/baselines/host-tabicl-a100.jsonl
new file mode 100644
index 0000000..f1fe321
--- /dev/null
+++ b/benchmarks/baselines/host-tabicl-a100.jsonl
@@ -0,0 +1,12 @@
+{"backend": "tabicl", "rows": 2000, "features": 500, "cells": 1000000, "cgroup_limit_gib": 64.0, "status": "ok", "fit_s": 6.3, "peak_growth_gib": 0.13, "retained_gib": 0.13, "table_gib": 0.004}
+{"backend": "tabicl", "rows": 5000, "features": 500, "cells": 2500000, "cgroup_limit_gib": 64.0, "status": "ok", "fit_s": 10.11, "peak_growth_gib": 0.2, "retained_gib": 0.13, "table_gib": 0.009}
+{"backend": "tabicl", "rows": 2000, "features": 1000, "cells": 2000000, "cgroup_limit_gib": 64.0, "status": "ok", "fit_s": 13.66, "peak_growth_gib": 0.24, "retained_gib": 0.21, "table_gib": 0.007}
+{"backend": "tabicl", "rows": 10000, "features": 500, "cells": 5000000, "cgroup_limit_gib": 64.0, "status": "ok", "fit_s": 15.08, "peak_growth_gib": 0.39, "retained_gib": 0.24, "table_gib": 0.019}
+{"backend": "tabicl", "rows": 100000, "features": 100, "cells": 10000000, "cgroup_limit_gib": 64.0, "status": "ok", "fit_s": 66.78, "peak_growth_gib": 45.11, "retained_gib": 42.8, "table_gib": 0.037}
+{"backend": "tabicl", "rows": 5000, "features": 1000, "cells": 5000000, "cgroup_limit_gib": 64.0, "status": "ok", "fit_s": 19.42, "peak_growth_gib": 0.39, "retained_gib": 0.24, "table_gib": 0.019}
+{"backend": "tabicl", "rows": 20000, "features": 500, "cells": 10000000, "cgroup_limit_gib": 64.0, "status": "ok", "fit_s": 58.84, "peak_growth_gib": 46.0, "retained_gib": 43.3, "table_gib": 0.037}
+{"backend": "tabicl", "rows": 30000, "features": 500, "cells": 15000000, "cgroup_limit_gib": 64.0, "status": "ok", "fit_s": 74.84, "peak_growth_gib": 46.17, "retained_gib": 43.37, "table_gib": 0.056}
+{"backend": "tabicl", "rows": 10000, "features": 500, "cells": 5000000, "cgroup_limit_gib": 64.0, "status": "ok", "fit_s": 14.89, "peak_growth_gib": 0.39, "retained_gib": 0.24, "table_gib": 0.019}
+{"backend": "tabicl", "rows": 20000, "features": 500, "cells": 10000000, "cgroup_limit_gib": 64.0, "status": "ok", "fit_s": 64.97, "peak_growth_gib": 46.0, "retained_gib": 0.29, "table_gib": 0.037}
+{"backend": "tabicl", "rows": 100000, "features": 100, "cells": 10000000, "cgroup_limit_gib": 64.0, "status": "ok", "fit_s": 73.18, "peak_growth_gib": 45.05, "retained_gib": 0.23, "table_gib": 0.037}
+{"backend": "tabicl", "rows": 30000, "features": 500, "cells": 15000000, "cgroup_limit_gib": 64.0, "status": "ok", "fit_s": 81.12, "peak_growth_gib": 46.17, "retained_gib": 0.37, "table_gib": 0.056}
diff --git a/benchmarks/bench_host_memory.py b/benchmarks/bench_host_memory.py
new file mode 100644
index 0000000..8f4e5fe
--- /dev/null
+++ b/benchmarks/bench_host_memory.py
@@ -0,0 +1,128 @@
+#!/usr/bin/env python3
+"""Host (CPU) memory per fit: the limit GPU admission doesn't see.
+
+With GPU admission bounded by what fits actually need (memory/
+requirement.py), the next ceiling is host RAM: the models' CPU-side
+preprocessing scales with cells, and a fit that exhausts the pod's memory
+gets the whole replica OOM-killed -- every tenant's context with it.
+
+For each shape this runs one fit in a fresh subprocess (an OOM kill takes
+only that child), sampling RSS every 20 ms, and records the peak RSS
+growth during fit and the RSS still held afterwards (host memory per
+cached context), plus the cgroup memory limit.
+
+Usage: bench_host_memory.py --backend tabicl --shapes 20000x100 10000x500 --out host.jsonl
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import multiprocessing as mp
+import sys
+import threading
+import time
+from pathlib import Path
+
+GiB = 1024**3
+
+
+def rss_bytes() -> int:
+    with open("/proc/self/status") as f:
+        for line in f:
+            if line.startswith("VmRSS:"):
+                return int(line.split()[1]) * 1024
+    return 0
+
+
+def cgroup_limit_bytes() -> int | None:
+    for path in (
+        "/sys/fs/cgroup/memory.max",
+        "/sys/fs/cgroup/memory/memory.limit_in_bytes",
+    ):
+        try:
+            with open(path) as f:
+                raw = f.read().strip()
+        except OSError:
+            continue
+        if raw.isdigit() and int(raw) < 1 << 60:
+            return int(raw)
+    return None
+
+
+def _child(backend_kind: str, device: str, rows: int, features: int, queue) -> None:
+    sys.path.insert(0, __file__.rsplit("/", 1)[0])
+    from profile_fit_predict import build_backend, make_table
+
+    backend = build_backend(backend_kind, device, "kv")
+    backend.warmup()
+    X, y = make_table(rows, features, "classification", seed=rows + features)
+    base = rss_bytes()
+    peak = [base]
+    stop = threading.Event()
+
+    def sample() -> None:
+        while not stop.is_set():
+            peak[0] = max(peak[0], rss_bytes())
+            time.sleep(0.02)
+
+    t = threading.Thread(target=sample, daemon=True)
+    t.start()
+    t0 = time.perf_counter()
+    payload = backend.fit(X, y, "classification")
+    fit_s = time.perf_counter() - t0
+    stop.set()
+    t.join()
+    del X, y
+    import gc
+
+    gc.collect()
+    queue.put(
+        {
+            "status": "ok",
+            "fit_s": round(fit_s, 2),
+            "peak_growth_gib": round((peak[0] - base) / GiB, 2),
+            "retained_gib": round((rss_bytes() - base) / GiB, 2),
+            "table_gib": round(rows * features * 4 / GiB, 3),
+        }
+    )
+    del payload
+
+
+def main() -> int:
+    p = argparse.ArgumentParser()
+    p.add_argument("--backend", choices=["tabicl", "tabpfn"], required=True)
+    p.add_argument("--device", default="cuda:0")
+    p.add_argument("--shapes", nargs="+", required=True)
+    p.add_argument("--out", default=None)
+    args = p.parse_args()
+    ctx = mp.get_context("spawn")
+    limit = cgroup_limit_bytes()
+    if args.out:
+        Path(args.out).parent.mkdir(parents=True, exist_ok=True)
+    out = open(args.out, "a") if args.out else None  # noqa: SIM115
+    for shape in args.shapes:
+        rows, features = (int(v) for v in shape.lower().split("x"))
+        q = ctx.Queue()
+        child = ctx.Process(
+            target=_child, args=(args.backend, args.device, rows, features, q)
+        )
+        child.start()
+        child.join()
+        rec = {
+            "backend": args.backend,
+            "rows": rows,
+            "features": features,
+            "cells": rows * features,
+        }
+        rec["cgroup_limit_gib"] = round(limit / GiB, 1) if limit else None
+        rec.update(q.get() if not q.empty() else {"status": f"exit {child.exitcode}"})
+        print(json.dumps(rec), flush=True)
+        if out:
+            out.write(json.dumps(rec) + "\n")
+            out.flush()
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/benchmarks/gke/devbox.sh b/benchmarks/gke/devbox.sh
index 84d7a61..0224105 100755
--- a/benchmarks/gke/devbox.sh
+++ b/benchmarks/gke/devbox.sh
@@ -91,6 +91,11 @@ up() {
   fi
   k wait --for=condition=Ready node -l cloud.google.com/gke-nodepool=gpu --timeout=20m
   local gpus; gpus=$(gpus_for "$MACHINE")
+  local phase; phase=$(k get pod $POD -o jsonpath='{.status.phase}' 2>/dev/null || true)
+  if [ -n "$phase" ] && [ "$phase" != "Running" ] && [ "$phase" != "Pending" ]; then
+    echo "==> pod is $phase (e.g. OOMKilled); recreating"
+    k delete pod $POD --wait
+  fi
   cat <<YAML | k apply -f -
 apiVersion: v1
 kind: Pod
@@ -111,7 +116,9 @@ spec:
     - name: box
       image: $IMAGE
       command: ["sleep", "infinity"]
-      resources: {limits: {nvidia.com/gpu: $gpus}}
+      # A memory limit, as KubeRay sets in prod: host-RAM exhaustion is then
+      # a cgroup OOM inside the container, and the limit is readable there.
+      resources: {limits: {nvidia.com/gpu: $gpus, memory: ${POD_MEMORY:-64Gi}}}
       env:
         - {name: TABPFN_MODEL_CACHE_DIR, value: /model-weights/tabctx/tabpfn}
         - {name: HF_HUB_CACHE, value: /model-weights/tabctx/hf/hub}
diff --git a/src/tabctx/backends/tabicl.py b/src/tabctx/backends/tabicl.py
index c170ad5..01240f9 100644
--- a/src/tabctx/backends/tabicl.py
+++ b/src/tabctx/backends/tabicl.py
@@ -213,6 +213,7 @@ def fit(self, X: ArrayLike, y: ArrayLike, task: Task) -> Any:
         # a future tabicl change must not silently re-pin (see
         # _unpin_shared_module for the OOM this prevents).
         self._unpin_shared_module(model)
+        self._release_host_memory(model)
         return model
 
     def _unpin_shared_module(self, model: Any) -> None:
@@ -231,6 +232,31 @@ def _unpin_shared_module(self, model: Any) -> None:
         if module is not None and hasattr(module, "clear_cache"):
             module.clear_cache()
 
+    def _release_host_memory(self, model: Any) -> None:
+        """Return the offload staging memory a large call leaves pinned.
+
+        Above ~5M cells tabicl offloads stage outputs to pinned host memory.
+        Each stage's InferenceManager keeps its PinnedBufferPool afterwards,
+        on the SHARED backbone, and the CUDA pinned-host allocator caches
+        what was freed. Measured on an A100 (20,000 x 500): +43 GiB of host
+        RSS after one fit, 10 GiB of it pooled buffers and 33 GiB allocator
+        cache, none of it the context's -- two large tables would exhaust a
+        64 GiB pod and get the replica OOM-killed. Clearing the pools and
+        emptying the host cache returns it (only unused cached blocks are
+        freed, so a concurrent call on another device is unaffected)."""
+        module = getattr(model, "model_", None)
+        if module is None:
+            return
+        for sub in module.modules():
+            pool = getattr(getattr(sub, "inference_mgr", None), "_buffer_pool", None)
+            if pool is not None:
+                pool.clear()
+        import torch
+
+        empty_host_cache = getattr(torch._C, "_host_emptyCache", None)
+        if callable(empty_host_cache):
+            empty_host_cache()
+
     def predict(
         self, payload: Any, X_test: ArrayLike, return_proba: bool = False
     ) -> PredictOutcome:
@@ -244,6 +270,7 @@ def predict(
             raise BackendComputeError(f"CUDA OOM during predict(): {e}") from e
         finally:
             self._unpin_shared_module(model)
+            self._release_host_memory(model)
 
     def _predict(
         self, model: Any, X_test: ArrayLike, return_proba: bool
diff --git a/src/tabctx/engine.py b/src/tabctx/engine.py
index d4cff88..bc423b4 100644
--- a/src/tabctx/engine.py
+++ b/src/tabctx/engine.py
@@ -6,6 +6,7 @@
 
 from __future__ import annotations
 
+import contextlib
 import uuid
 
 from tabctx.backends.base import TabularICLBackend
@@ -13,6 +14,7 @@
 from tabctx.chunking import choose_chunk_rows, split_rows
 from tabctx.errors import AdmissionRejected, DatasetNotFoundError, InvalidInputError
 from tabctx.memory.estimator import MemoryEstimator
+from tabctx.memory.host import HostMemoryGuard
 from tabctx.types import ArrayLike, EngineStats, PredictOutcome, Task
 
 
@@ -99,6 +101,7 @@ def __init__(
         backends: dict[str, TabularICLBackend] | None = None,
         estimators: dict[str, MemoryEstimator] | None = None,
         default_backend: str | None = None,
+        host_guard: HostMemoryGuard | None = None,
     ) -> None:
         """Single-backend form: TabctxEngine(backend=, cache=, estimator=).
         Multi-backend form: TabctxEngine(backends={name: b}, cache=,
@@ -122,6 +125,9 @@ def __init__(
         self._estimators = estimators
         self._default_backend = default_backend
         self._cache = cache
+        # Process-wide host-memory admission, shared by every device's
+        # engine (see memory/host.py); None = unguarded.
+        self._host_guard = host_guard
 
     @property
     def backend_names(self) -> list[str]:
@@ -211,7 +217,14 @@ def fit(
         estimator = self._estimators[backend_name]
 
         resolved_id = dataset_id or str(uuid.uuid4())
-        with self._cache.lock:
+        # Host memory first, outside the cache lock: a large fit waiting for
+        # a host slot must not block predicts on this device meanwhile.
+        host_slot = (
+            self._host_guard.reserve(n_train, n_features)
+            if self._host_guard is not None
+            else contextlib.nullcontext()
+        )
+        with host_slot, self._cache.lock:
             # Usage-aware admission (v0.9.0): what OOMs is the fit's
             # transient PEAK on top of everything resident, so the
             # estimated peak must fit the headroom given current cache
diff --git a/src/tabctx/memory/host.py b/src/tabctx/memory/host.py
new file mode 100644
index 0000000..e54b5b1
--- /dev/null
+++ b/src/tabctx/memory/host.py
@@ -0,0 +1,95 @@
+"""Host (CPU) memory admission: the limit GPU admission doesn't see.
+
+Large TabICL fits offload stage outputs to pinned host memory while they
+run. Measured on an A100 pod (benchmarks/bench_host_memory.py,
+benchmarks/baselines/host-tabicl-a100.jsonl): fits up to ~5M cells grow
+host RSS by < 0.5 GiB, but 10-15M-cell fits (20,000 x 500, 100,000 x 100)
+peak at ~46 GiB of host memory -- transient, returned afterwards (see
+TabICLBackend._release_host_memory). One such fit fits a 64 GiB pod; two
+at once, on two GPUs of one replica, get it OOM-killed with every tenant's
+context.
+
+HostMemoryGuard lets at most budget // large_fit_bytes large fits run at
+once, process-wide (shared by every device's engine). A large fit waits
+for a slot instead of exhausting the pod; small fits never wait. "Large"
+is a cell count below which no offload was ever measured.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import threading
+from collections.abc import Iterator
+from pathlib import Path
+
+GiB = 1024**3
+# Calibrated on an A100-40GB (see module docstring).
+LARGE_FIT_CELLS = 3_000_000
+LARGE_FIT_HOST_BYTES = 48 * GiB
+
+
+def detect_host_budget_bytes(fraction: float = 0.85) -> int | None:
+    """This process's memory limit (cgroup v2/v1, else physical RAM) times
+    fraction, or None when it can't be read."""
+    for path in (
+        "/sys/fs/cgroup/memory.max",
+        "/sys/fs/cgroup/memory/memory.limit_in_bytes",
+    ):
+        try:
+            raw = Path(path).read_text().strip()
+        except OSError:
+            continue
+        if raw.isdigit() and int(raw) < 1 << 60:
+            return int(int(raw) * fraction)
+    try:
+        for line in Path("/proc/meminfo").read_text().splitlines():
+            if line.startswith("MemTotal:"):
+                return int(int(line.split()[1]) * 1024 * fraction)
+    except OSError:
+        pass
+    return None
+
+
+class HostMemoryGuard:
+    def __init__(
+        self,
+        budget_bytes: int,
+        large_fit_cells: int = LARGE_FIT_CELLS,
+        large_fit_bytes: int = LARGE_FIT_HOST_BYTES,
+    ) -> None:
+        self.slots = max(1, budget_bytes // large_fit_bytes)
+        self._large_fit_cells = large_fit_cells
+        self._semaphore = threading.BoundedSemaphore(self.slots)
+        self._waiting = 0
+        self._active = 0
+        self._lock = threading.Lock()
+
+    def is_large(self, n_train: int, n_features: int) -> bool:
+        return n_train * n_features >= self._large_fit_cells
+
+    @contextlib.contextmanager
+    def reserve(self, n_train: int, n_features: int) -> Iterator[None]:
+        """Hold a large-fit slot for the duration (no-op for small fits)."""
+        if not self.is_large(n_train, n_features):
+            yield
+            return
+        with self._lock:
+            self._waiting += 1
+        self._semaphore.acquire()
+        with self._lock:
+            self._waiting -= 1
+            self._active += 1
+        try:
+            yield
+        finally:
+            with self._lock:
+                self._active -= 1
+            self._semaphore.release()
+
+    def stats(self) -> dict[str, int]:
+        with self._lock:
+            return {
+                "large_fit_slots": self.slots,
+                "large_fits_active": self._active,
+                "large_fits_waiting": self._waiting,
+            }
diff --git a/src/tabctx/serve/app.py b/src/tabctx/serve/app.py
index 982f66a..a79a538 100644
--- a/src/tabctx/serve/app.py
+++ b/src/tabctx/serve/app.py
@@ -216,6 +216,7 @@ def __init__(self) -> None:
         self._backend_name = built.backend.name  # default model id
         self._spill_store = built.spill_store
         self._spill_stores = built.spill_stores
+        self._host_guard = built.host_guard
         self._devices = built.devices
         self._device = built.device
         self._gpu_capacity_bytes = built.gpu_capacity_bytes
@@ -608,6 +609,7 @@ def readyz(self):
             "spill": self._spill_stats(),
             "real_gpu_memory": real_gpu_memory,
             "devices": self._device_breakdown(),
+            "host_memory": self._host_guard.stats() if self._host_guard else None,
         }
 
     def _cuda_devices(self) -> list[str]:
diff --git a/src/tabctx/serve/factory.py b/src/tabctx/serve/factory.py
index a40cd9a..041680e 100644
--- a/src/tabctx/serve/factory.py
+++ b/src/tabctx/serve/factory.py
@@ -80,6 +80,7 @@
     DEFAULT_HARD_CEILING_BYTES,
     MemoryEstimator,
 )
+from tabctx.memory.host import HostMemoryGuard, detect_host_budget_bytes
 from tabctx.memory.requirement import ResidentRateModel, WorkingFloorModel
 from tabctx.multidevice import MultiDeviceEngine
 from tabctx.serve.uploads import DEFAULT_MAX_UPLOAD_BYTES, DEFAULT_TTL_S
@@ -247,6 +248,7 @@ class BuiltEngine:
     devices: tuple[str, ...] = ()
     # One per device when there are several (spill_store is then None).
     spill_stores: tuple = ()
+    host_guard: HostMemoryGuard | None = None
 
     # Single-model conveniences (the common deployment):
     @property
@@ -502,6 +504,11 @@ def build_engine(settings: ServeSettings | None = None) -> BuiltEngine:
     if settings.spill_dir:
         _warn_on_spill_layout(settings.spill_dir, multi)
     estimators: dict[str, MemoryEstimator] = {}
+    # One host-memory guard per process: host RAM is shared by every device.
+    host_budget = (
+        None if set(settings.backends) == {"fake"} else detect_host_budget_bytes()
+    )
+    host_guard = HostMemoryGuard(host_budget) if host_budget else None
     shards: list[TabctxEngine] = []
     spill_stores: list[DiskSpillStore | None] = []
     first_backends: dict[str, TabularICLBackend] = {}
@@ -536,6 +543,7 @@ def build_engine(settings: ServeSettings | None = None) -> BuiltEngine:
                 estimators={name: estimators[name] for name in backends},
                 cache=cache,
                 default_backend=default,
+                host_guard=host_guard,
             )
         )
         spill_stores.append(spill_store)
@@ -554,4 +562,5 @@ def build_engine(settings: ServeSettings | None = None) -> BuiltEngine:
         gpu_capacity_bytes=gpu_capacity,
         devices=devices,
         spill_stores=tuple(s for s in spill_stores if s is not None) if multi else (),
+        host_guard=host_guard,
     )
diff --git a/tests/unit/test_host_memory.py b/tests/unit/test_host_memory.py
new file mode 100644
index 0000000..c056632
--- /dev/null
+++ b/tests/unit/test_host_memory.py
@@ -0,0 +1,99 @@
+"""Unit tests for process-wide host-memory admission (memory/host.py)."""
+
+import threading
+
+from tabctx.backends.fake import FakeBackend
+from tabctx.cache.manager import ContextCacheManager
+from tabctx.engine import TabctxEngine
+from tabctx.memory import (
+    A100_40GB_TABICL_CALIBRATION,
+    AdaptiveMemoryEstimator,
+    PowerLawMemoryEstimator,
+)
+from tabctx.memory.host import GiB, HostMemoryGuard
+
+
+def test_slots_from_budget():
+    assert HostMemoryGuard(176 * GiB, large_fit_bytes=48 * GiB).slots == 3
+    assert HostMemoryGuard(10 * GiB, large_fit_bytes=48 * GiB).slots == 1  # never 0
+
+
+def test_small_fits_never_wait():
+    guard = HostMemoryGuard(GiB, large_fit_cells=100, large_fit_bytes=GiB)
+    with guard.reserve(50, 1), guard.reserve(50, 1):  # both small: no slot used
+        assert guard.stats()["large_fits_active"] == 0
+
+
+def test_large_fits_are_limited_to_the_slots():
+    guard = HostMemoryGuard(
+        2 * GiB, large_fit_cells=100, large_fit_bytes=GiB
+    )  # 2 slots
+    inside, release = threading.Semaphore(0), threading.Event()
+
+    def large():
+        with guard.reserve(100, 1):
+            inside.release()
+            release.wait(5)
+
+    threads = [threading.Thread(target=large, daemon=True) for _ in range(3)]
+    for t in threads:
+        t.start()
+    assert inside.acquire(timeout=5) and inside.acquire(timeout=5)
+    assert not inside.acquire(timeout=0.3)  # the third waits for a slot
+    assert guard.stats() == {
+        "large_fit_slots": 2,
+        "large_fits_active": 2,
+        "large_fits_waiting": 1,
+    }
+    release.set()
+    for t in threads:
+        t.join(5)
+    assert guard.stats()["large_fits_active"] == 0
+
+
+class GatedFit(FakeBackend):
+    def __init__(self) -> None:
+        super().__init__(bytes_hint=10)
+        self.gate = threading.Event()
+        self.entered = threading.Event()
+
+    def fit(self, X, y, task):
+        if len(X) >= 100:
+            self.entered.set()
+            assert self.gate.wait(5)
+        return super().fit(X, y, task)
+
+
+def test_a_waiting_large_fit_does_not_block_predicts_on_its_device():
+    guard = HostMemoryGuard(GiB, large_fit_cells=100, large_fit_bytes=GiB)  # 1 slot
+    backends = [GatedFit(), GatedFit()]
+    est = AdaptiveMemoryEstimator(
+        fallback=PowerLawMemoryEstimator(A100_40GB_TABICL_CALIBRATION)
+    )
+    engines = [
+        TabctxEngine(
+            backend=b,
+            cache=ContextCacheManager(est.ceiling_bytes()),
+            estimator=est,
+            host_guard=guard,
+        )
+        for b in backends
+    ]
+    engines[1].fit([[1.0]] * 3, ["a"] * 3, dataset_id="small")
+    big = ([[1.0]] * 100, ["a"] * 100)
+    t0 = threading.Thread(target=engines[0].fit, args=big, daemon=True)
+    t0.start()  # takes the only slot, then blocks inside the backend
+    assert backends[0].entered.wait(5)
+    t1 = threading.Thread(target=engines[1].fit, args=big, daemon=True)
+    t1.start()  # waits for the host slot
+    # Device 1 still serves predicts while its large fit waits for host memory.
+    done = threading.Event()
+    threading.Thread(
+        target=lambda: (engines[1].predict("small", [[1.0]]), done.set()), daemon=True
+    ).start()
+    assert done.wait(3), "a fit waiting on host memory blocked its device"
+    assert guard.stats()["large_fits_waiting"] == 1
+    backends[0].gate.set()
+    backends[1].gate.set()
+    t0.join(5)
+    t1.join(5)

From 880959644f092a6d94a97c20322fc396819b7317 Mon Sep 17 00:00:00 2001
From: Amrit Krishnan <amrit110@gmail.com>
Date: Sat, 17 Oct 2026 16:39:03 -0400
Subject: [PATCH 06/39] fix(tabicl): give tabicl a disk offload dir; leave host
 headroom for its staging

Without disk_offload_dir, an offloaded output larger than host memory
fell through to cpu (swap): an OOM kill in a swapless pod (65000x500 on a
64 GiB A100 pod). TABCTX_OFFLOAD_DIR (default <spill>/offload, per device)
lets it degrade to disk memmaps; cpu_safety_factor 0.6 leaves room for
its pinned staging buffers and other devices' work.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 benchmarks/sweep_engine.py       |  3 ++-
 src/tabctx/backends/tabicl.py    | 25 +++++++++++++++++++++++--
 src/tabctx/serve/factory.py      | 28 +++++++++++++++++++++++++++-
 tests/unit/test_serve_factory.py |  1 +
 4 files changed, 53 insertions(+), 4 deletions(-)

diff --git a/benchmarks/sweep_engine.py b/benchmarks/sweep_engine.py
index 81c01f5..354527c 100644
--- a/benchmarks/sweep_engine.py
+++ b/benchmarks/sweep_engine.py
@@ -31,6 +31,7 @@ def main() -> int:
     p.add_argument("--backend", default="tabicl")
     p.add_argument("--features", type=int, nargs="+", default=[10, 100, 500, 1000])
     p.add_argument("--peak-only", action="store_true")
+    p.add_argument("--min-rows", type=int, default=0)
     p.add_argument("--out", default=None)
     args = p.parse_args()
     if args.peak_only:
@@ -41,7 +42,7 @@ def main() -> int:
     out = open(args.out, "a") if args.out else None  # noqa: SIM115
     policy = "peak-only" if args.peak_only else "requirement"
     for d in args.features:
-        for n in ROWS:
+        for n in (r for r in ROWS if r >= args.min_rows):
             X = rng.normal(size=(n, d)).astype(np.float32)
             y = (X[:, 0] > 0).astype(int)
             rec = {"policy": policy, "rows": n, "features": d}
diff --git a/src/tabctx/backends/tabicl.py b/src/tabctx/backends/tabicl.py
index 01240f9..a92c478 100644
--- a/src/tabctx/backends/tabicl.py
+++ b/src/tabctx/backends/tabicl.py
@@ -96,6 +96,8 @@ def __init__(
         device: str | None = None,
         kv_cache: bool | str = "kv",
         kv_precision: KVPrecision = "native",
+        offload_dir: str | None = None,
+        cpu_safety_factor: float = 0.6,
     ) -> None:
         """kv_cache: passed through to TabICL ("kv", "repr", True, or
         False). tabicl's own default is False, which makes every
@@ -117,6 +119,16 @@ def __init__(
             raise ValueError(f"kv_precision must be one of {KV_PRECISIONS}")
         # Storage of the ICL key/value cache (see backends/tabicl_kv.py).
         self._kv_precision: KVPrecision = kv_precision
+        # Large fits offload stage outputs off the GPU (tabicl's
+        # offload="auto"). Without a disk directory, an output too big for
+        # host memory falls through to "cpu (swap)" -- an OOM kill in a
+        # container with no swap (measured: 65,000 x 500 on a 64 GiB pod).
+        # With one, it degrades to disk memmaps instead. tabicl budgets host
+        # memory as cgroup headroom x cpu_safety_factor but doesn't count
+        # its own pinned staging buffers or other devices' concurrent work,
+        # hence 0.6 rather than its 0.85 default.
+        self._offload_dir = offload_dir
+        self._offload_config = {"cpu_safety_factor": cpu_safety_factor}
         self._last_context_bytes: int | None = None
         self._last_fit_peak_bytes: int | None = None
         # task -> attrs of an already-loaded backbone, shared across
@@ -136,14 +148,23 @@ def __init__(
         _install_restore_hook(TabICLRegressor)
 
     def _make_estimator(self, task: Task) -> Any:
+        options = {
+            "device": self._device,
+            "kv_cache": self._kv_cache,
+            "disk_offload_dir": self._offload_dir,
+            "inference_config": {
+                stage: dict(self._offload_config)
+                for stage in ("COL_CONFIG", "ROW_CONFIG", "ICL_CONFIG")
+            },
+        }
         if task == "classification":
             from tabicl import TabICLClassifier
 
-            est = TabICLClassifier(device=self._device, kv_cache=self._kv_cache)
+            est = TabICLClassifier(**options)
         else:
             from tabicl import TabICLRegressor
 
-            est = TabICLRegressor(device=self._device, kv_cache=self._kv_cache)
+            est = TabICLRegressor(**options)
 
         shared = self._shared_backbones.get(task)
         if shared is not None:
diff --git a/src/tabctx/serve/factory.py b/src/tabctx/serve/factory.py
index 041680e..bbadbc9 100644
--- a/src/tabctx/serve/factory.py
+++ b/src/tabctx/serve/factory.py
@@ -36,6 +36,10 @@
   (benchmarks/baselines/kv-precision-tabicl-a100.jsonl): identical labels
   and metrics, max probability drift 0.009, at up to ~1.4x warm predict
   latency on large contexts. TabPFN already stores its cache in int8.
+- ``TABCTX_OFFLOAD_DIR``: where TabICL may spill stage outputs too large
+  for GPU and host memory (memory-mapped files), instead of exhausting the
+  pod's RAM; default ``<TABCTX_SPILL_DIR>/offload``, else a temporary
+  directory. One subdirectory per device.
 - ``TABCTX_BATCH_WINDOW_MS``: coalescing window for same-context predict
   batching (see batching.py); default 5. 0 disables coalescing.
 - ``TABCTX_MAX_UPLOAD_BYTES`` (default 4GiB) and ``TABCTX_UPLOAD_TTL_S``
@@ -61,6 +65,7 @@
 
 import logging
 import os
+import tempfile
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Literal
@@ -89,6 +94,7 @@
 GPU_MEMORY_FRACTION_ENV_VAR = "TABCTX_GPU_MEMORY_FRACTION"
 KV_CACHE_ENV_VAR = "TABCTX_KV_CACHE"
 KV_PRECISION_ENV_VAR = "TABCTX_KV_PRECISION"
+OFFLOAD_DIR_ENV_VAR = "TABCTX_OFFLOAD_DIR"
 BATCH_WINDOW_MS_ENV_VAR = "TABCTX_BATCH_WINDOW_MS"
 MAX_UPLOAD_BYTES_ENV_VAR = "TABCTX_MAX_UPLOAD_BYTES"
 SPILL_DIR_ENV_VAR = "TABCTX_SPILL_DIR"
@@ -117,6 +123,8 @@ class ServeSettings:
     spill_capacity_bytes: int = DEFAULT_SPILL_CAPACITY_BYTES
     # Empty = "auto" (see resolve_devices).
     devices: tuple[str, ...] = ()
+    # None = <spill_dir>/offload, else a temporary directory.
+    offload_dir: str | None = None
 
     @classmethod
     def from_env(cls) -> ServeSettings:
@@ -192,6 +200,7 @@ def from_env(cls) -> ServeSettings:
         if spill_capacity <= 0:
             raise ValueError(f"{SPILL_CAPACITY_ENV_VAR} must be positive")
         devices = _parse_devices(os.environ.get(DEVICES_ENV_VAR, "auto"))
+        offload_dir = os.environ.get(OFFLOAD_DIR_ENV_VAR) or None
         return cls(
             backends=backends,
             gpu_memory_fraction=fraction,
@@ -203,6 +212,7 @@ def from_env(cls) -> ServeSettings:
             spill_dir=spill_dir,
             spill_capacity_bytes=spill_capacity,
             devices=devices,
+            offload_dir=offload_dir,
         )
 
 
@@ -424,7 +434,10 @@ def _build_real_backend(
             False if settings.kv_cache == "off" else settings.kv_cache
         )
         backend = TabICLBackend(
-            device=device, kv_cache=kv_cache, kv_precision=settings.kv_precision
+            device=device,
+            kv_cache=kv_cache,
+            kv_precision=settings.kv_precision,
+            offload_dir=_offload_dir(settings, device),
         )
     # Weights load at replica start, not inside (and charged to) the first
     # tenant's fit.
@@ -432,6 +445,19 @@ def _build_real_backend(
     return backend
 
 
+def _offload_dir(settings: ServeSettings, device: str) -> str:
+    """Per-device directory for TabICL's disk offload (created)."""
+    if settings.offload_dir:
+        root = settings.offload_dir
+    elif settings.spill_dir:
+        root = os.path.join(settings.spill_dir, "offload")
+    else:
+        root = os.path.join(tempfile.gettempdir(), "tabctx-offload")
+    path = os.path.join(root, device.replace(":", "_"))
+    os.makedirs(path, exist_ok=True)
+    return path
+
+
 def _build_backend(
     kind: BackendKind, settings: ServeSettings, device: str = "cpu"
 ) -> TabularICLBackend:
diff --git a/tests/unit/test_serve_factory.py b/tests/unit/test_serve_factory.py
index 0cc0939..7f8372f 100644
--- a/tests/unit/test_serve_factory.py
+++ b/tests/unit/test_serve_factory.py
@@ -253,6 +253,7 @@ def test_settings_defaults_match_from_env_defaults(self, monkeypatch):
             SPILL_CAPACITY_ENV_VAR,
             DEVICES_ENV_VAR,
             "TABCTX_KV_PRECISION",
+            "TABCTX_OFFLOAD_DIR",
         ):
             monkeypatch.delenv(var, raising=False)
         assert ServeSettings.from_env() == ServeSettings()

From 4a503ac147583a99b917ec3e91057422733fab18 Mon Sep 17 00:00:00 2001
From: Amrit Krishnan <amrit110@gmail.com>
Date: Sat, 17 Oct 2026 16:46:44 -0400
Subject: [PATCH 07/39] fix: review findings on #40

1. Admission learns the fit-time (pre-compression) resident: the fit
   builds the native kv cache before int8 compression shrinks it
   (fit_resident_bytes_hint); compressed size stays cache accounting.
2. TabICL fit releases pinned host memory in a finally (a failed large
   fit kept ~43 GiB pinned while its host slot was released).
3. Host guard follows tabicl actual offload trigger (may_offload: output
   > half of currently free GPU memory) and guards predicts too; each
   slot budgets 0.6/slots of host headroom; budget minus baseline RSS;
   warns when below one large fit.
4. Requirement grid uses measured min_free uncapped; record_fit no
   longer caps at a peak measured under pressure.
5. int8 dequantize into one fp16 buffer; chunked packing; engine splits
   an OOM predict chunk until it fits.
6. Admission-effect test now depends on the floor model; tests for
   ResidentRateModel, extrapolation, factory wiring, host budget
   detection, release-on-failure, backend offload hint; real-tabicl
   CPU integration test for compressed KV (fit/predict/spill).
7. Floor model monotone in rows; warning if private release hooks vanish;
   docs corrected.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 benchmarks/gen_requirement_grid.py            |  10 +-
 src/tabctx/backends/tabicl.py                 |  90 +++++++++----
 src/tabctx/backends/tabicl_kv.py              |  40 ++++--
 src/tabctx/engine.py                          |  76 ++++++++---
 src/tabctx/memory/adaptive.py                 |  13 +-
 .../calibration_tabicl_a100_requirement.py    |  24 ++--
 src/tabctx/memory/host.py                     |  84 +++++++++---
 src/tabctx/memory/requirement.py              |  20 ++-
 src/tabctx/serve/factory.py                   |  13 +-
 tests/integration/test_tabicl_kv_real.py      |  55 ++++++++
 tests/unit/test_engine.py                     |  69 ++++++++++
 tests/unit/test_host_memory.py                | 102 +++++++++++++-
 tests/unit/test_requirement.py                | 126 ++++++++++++++++--
 tests/unit/test_serve_factory.py              |  12 +-
 tests/unit/test_tabicl_host_release.py        |  70 ++++++++++
 15 files changed, 684 insertions(+), 120 deletions(-)
 create mode 100644 tests/integration/test_tabicl_kv_real.py
 create mode 100644 tests/unit/test_tabicl_host_release.py

diff --git a/benchmarks/gen_requirement_grid.py b/benchmarks/gen_requirement_grid.py
index 26a3637..724463f 100644
--- a/benchmarks/gen_requirement_grid.py
+++ b/benchmarks/gen_requirement_grid.py
@@ -41,11 +41,11 @@ def main() -> int:
     for r in ok:
         n, d = r["rows"], r["features"]
         resident = int(r["resident_gib"] * GiB)
-        # The floor can't exceed what the fit used unconstrained: a search
-        # result above that is ballast fragmentation, not need.
-        used = int((r["unconstrained_peak_gib"] - r["resident_gib"]) * GiB)
-        working = min(int(r["min_working_gib"] * GiB), used)
-        need = resident + working
+        # Uncapped: min_free is what a live device needs (it includes
+        # allocator fragmentation and workspaces that max_memory_allocated
+        # doesn't see), even where it exceeds the unconstrained peak.
+        working = int(r["min_working_gib"] * GiB)
+        need = int(r["min_free_gib"] * GiB)
         note = (
             f"resident {r['resident_gib']} GiB, unconstrained peak "
             f"{r['unconstrained_peak_gib']} GiB, fit {r['unconstrained_fit_s']}s "
diff --git a/src/tabctx/backends/tabicl.py b/src/tabctx/backends/tabicl.py
index a92c478..bd80fd0 100644
--- a/src/tabctx/backends/tabicl.py
+++ b/src/tabctx/backends/tabicl.py
@@ -49,6 +49,7 @@
 
 from __future__ import annotations
 
+import logging
 import threading
 from typing import Any
 
@@ -63,6 +64,10 @@
 # calls _load_model(), which would otherwise torch.load the full checkpoint
 # and move a second backbone onto the GPU on every spill restore.
 _restoring = threading.local()
+_warned = [False]  # one missing-hook warning per process
+# Offloaded stage-output bytes per cell (see TabICLBackend.may_offload).
+_OUTPUT_BYTES_PER_CELL = 2560
+logger = logging.getLogger(__name__)
 _TASK_TAGS = {"classification": b"C", "regression": b"R"}
 
 
@@ -131,6 +136,7 @@ def __init__(
         self._offload_config = {"cpu_safety_factor": cpu_safety_factor}
         self._last_context_bytes: int | None = None
         self._last_fit_peak_bytes: int | None = None
+        self._last_fit_resident_bytes: int | None = None
         # task -> attrs of an already-loaded backbone, shared across
         # fits. Loading is ~hundreds of MB of torch.load + H2D transfer
         # PER FIT without this. Sharing one nn.Module across estimators
@@ -201,40 +207,43 @@ def fit(self, X: ArrayLike, y: ArrayLike, task: Task) -> Any:
         before_bytes = torch.cuda.memory_allocated(device) if measuring else 0
         if measuring:
             torch.cuda.reset_peak_memory_stats(device)
+        self._last_context_bytes = None
+        self._last_fit_peak_bytes = None
+        self._last_fit_resident_bytes = None
         try:
             with device_scope(device):
                 model.fit(X, y)
-                # Before measuring: the context's resident size is the
-                # compressed one.
+                # The context as the fit built it, at full precision: what
+                # the fit needed resident, and so what admission must
+                # learn -- compression only shrinks it afterwards.
+                fit_resident = (
+                    torch.cuda.memory_allocated(device) - before_bytes
+                    if measuring
+                    else None
+                )
                 compress_icl_cache(model, self._kv_precision)
         except torch.cuda.OutOfMemoryError as e:
-            self._last_context_bytes = None
-            self._last_fit_peak_bytes = None
             raise BackendComputeError(f"CUDA OOM during fit(): {e}") from e
-        # Two DIFFERENT real measurements, for two different consumers --
-        # conflating them was the v0.4.0 design flaw fixed in v0.9.0:
-        # - resident delta (context_bytes_hint): what the fitted context
-        #   occupies afterward -> cache capacity/eviction accounting.
-        # - peak delta (fit_peak_bytes_hint): the transient high-water
-        #   during fit -> what ADMISSION must bound, since this is what
-        #   actually OOMs. Peak can exceed resident by orders of
-        #   magnitude (activations, ensemble intermediates).
+        finally:
+            # Even on failure: a large fit that dies after offloading leaves
+            # tens of GiB pinned, and its host slot is released right after.
+            self._unpin_shared_module(model)
+            self._release_host_memory(model)
+        # Three real measurements, for different consumers:
+        # - resident (context_bytes_hint): what the fitted context occupies
+        #   afterwards, compressed -> cache capacity/eviction accounting.
+        # - fit-time resident (fit_resident_bytes_hint): the context before
+        #   compression -> what admission learns the fit needed resident.
+        # - peak (fit_peak_bytes_hint): the transient high-water during fit.
         # None on CPU; the engine falls back to the formula estimator.
         if measuring:
             after_bytes = torch.cuda.memory_allocated(device)
             self._last_context_bytes = max(0, after_bytes - before_bytes)
+            self._last_fit_resident_bytes = max(0, fit_resident or 0)
             self._last_fit_peak_bytes = max(
                 0, torch.cuda.max_memory_allocated(device) - before_bytes
             )
-        else:
-            self._last_context_bytes = None
-            self._last_fit_peak_bytes = None
         self._stash_backbone(task, model)
-        # Explicit even though tabicl's fit path clears internally today:
-        # a future tabicl change must not silently re-pin (see
-        # _unpin_shared_module for the OOM this prevents).
-        self._unpin_shared_module(model)
-        self._release_host_memory(model)
         return model
 
     def _unpin_shared_module(self, model: Any) -> None:
@@ -253,6 +262,20 @@ def _unpin_shared_module(self, model: Any) -> None:
         if module is not None and hasattr(module, "clear_cache"):
             module.clear_cache()
 
+    def may_offload(self, n_rows: int, n_features: int) -> bool:
+        """Whether a call over n_rows x n_features may offload stage outputs
+        to host memory. Mirrors tabicl's auto-offload trigger -- an output
+        above auto_offload_threshold (0.5) x the device's CURRENTLY free GPU
+        memory -- with the output size calibrated on an A100 (no offload at
+        5M cells, offload at 10M, ~38 GB free): ~2.5 KB per cell for 8
+        estimators. So on a busy card, smaller tables offload too."""
+        if not is_cuda(self._device):
+            return False
+        import torch
+
+        free, _ = torch.cuda.mem_get_info(torch.device(self._device))
+        return _OUTPUT_BYTES_PER_CELL * n_rows * n_features > 0.4 * free
+
     def _release_host_memory(self, model: Any) -> None:
         """Return the offload staging memory a large call leaves pinned.
 
@@ -264,12 +287,16 @@ def _release_host_memory(self, model: Any) -> None:
         cache, none of it the context's -- two large tables would exhaust a
         64 GiB pod and get the replica OOM-killed. Clearing the pools and
         emptying the host cache returns it (only unused cached blocks are
-        freed, so a concurrent call on another device is unaffected)."""
+        freed; cudaFreeHost synchronizes, so a concurrent call on another
+        device may stall briefly, but its buffers are untouched)."""
         module = getattr(model, "model_", None)
         if module is None:
             return
-        for sub in module.modules():
-            pool = getattr(getattr(sub, "inference_mgr", None), "_buffer_pool", None)
+        pools = [
+            getattr(getattr(sub, "inference_mgr", None), "_buffer_pool", None)
+            for sub in module.modules()
+        ]
+        for pool in pools:
             if pool is not None:
                 pool.clear()
         import torch
@@ -277,6 +304,17 @@ def _release_host_memory(self, model: Any) -> None:
         empty_host_cache = getattr(torch._C, "_host_emptyCache", None)
         if callable(empty_host_cache):
             empty_host_cache()
+        if (not any(pools) or not callable(empty_host_cache)) and not _warned[0]:
+            # Private hooks (tabicl's pools, torch's host cache): a rename
+            # would silently bring back tens of GiB of retained host memory.
+            _warned[0] = True
+            logger.warning(
+                "tabctx could not release TabICL's pinned host memory "
+                "(buffer pools found: %s, torch host cache API: %s); large "
+                "fits may retain host RAM -- check tabicl/torch versions",
+                any(pools),
+                callable(empty_host_cache),
+            )
 
     def predict(
         self, payload: Any, X_test: ArrayLike, return_proba: bool = False
@@ -332,6 +370,12 @@ def context_bytes_hint(self, n_train: int, n_features: int) -> int | None:
         # the engine falls back to the formula-based MemoryEstimator.
         return self._last_context_bytes
 
+    def fit_resident_bytes_hint(self) -> int | None:
+        """The most recent fit's context before compression -- what it
+        needed resident while it ran (== context_bytes_hint at native
+        precision); None on CPU/after OOM."""
+        return self._last_fit_resident_bytes
+
     def fit_peak_bytes_hint(self) -> int | None:
         """Transient high-water memory of the most recent fit() -- the
         admission-relevant quantity (see fit()); None on CPU/after OOM."""
diff --git a/src/tabctx/backends/tabicl_kv.py b/src/tabctx/backends/tabicl_kv.py
index 2fe46a5..6e5eb03 100644
--- a/src/tabctx/backends/tabicl_kv.py
+++ b/src/tabctx/backends/tabicl_kv.py
@@ -1,10 +1,12 @@
 """Compressed storage for TabICL's ICL key/value cache.
 
-A TabICL context in kv mode is dominated by the ICL transformer's cached
-keys and values: ~0.19 MiB per training row for the default 8 estimators
-(measured on an A100: 100k x 10 = 17.6 GiB resident), stored in fp16 once
-AMP is on (>= 1024 rows) and fp32 below that. That is what bounds the
-longest table a GPU can hold, and how many contexts fit.
+For long tables a TabICL kv context is dominated by the ICL transformer's
+cached keys and values: ~0.19 MiB per training row for the default 8
+estimators (measured on an A100: 100k x 10 = 17.6 GiB resident), stored in
+fp16 once AMP is on (>= 1024 rows) and fp32 below that. That is what
+bounds the longest table a GPU can hold, and how many contexts fit. (Wide,
+short tables are dominated by the column cache instead -- ~0.37 of the
+context at 2,000 x 1,000 is ICL cache -- so compression saves less there.)
 
 This module swaps the ICL cache's entries for ``CompressedKVEntry``, which
 stores them compactly and decompresses on access:
@@ -51,17 +53,31 @@ def __init__(self, data: Tensor, scale: Tensor | None, dtype: torch.dtype) -> No
 
     @classmethod
     def pack(cls, t: Tensor, precision: KVPrecision) -> _Packed:
-        if precision == "int8":
-            absmax = t.abs().amax(dim=-1, keepdim=True).float().clamp_min(1e-8)
-            scale = (absmax / 127.0).to(torch.float16)
-            q = torch.round(t.float() / scale.float()).clamp_(-127, 127).to(torch.int8)
-            return cls(q, scale, t.dtype)
-        return cls(t.to(torch.float16), None, t.dtype)
+        if precision != "int8":
+            return cls(t.to(torch.float16), None, t.dtype)
+        q = torch.empty(t.shape, dtype=torch.int8, device=t.device)
+        scale = torch.empty((*t.shape[:-1], 1), dtype=torch.float16, device=t.device)
+        # Chunked along the batch dimension: the float32 working copy is one
+        # chunk, not the whole cache, on top of a context that is already
+        # resident at full precision.
+        step = max(1, t.shape[0] // 8) if t.dim() > 1 else t.shape[0]
+        for i in range(0, t.shape[0], step):
+            x = t[i : i + step].float()
+            s = (x.abs().amax(dim=-1, keepdim=True).clamp_min(1e-8) / 127.0).to(
+                torch.float16
+            )
+            scale[i : i + step] = s
+            q[i : i + step] = torch.round(x / s.float()).clamp_(-127, 127)
+        return cls(q, scale, t.dtype)
 
     def unpack(self) -> Tensor:
         if self.scale is None:
             return self.data.to(self.dtype)
-        return (self.data.to(torch.float16) * self.scale).to(self.dtype)
+        # One fp16 buffer, no int8->fp16 temporary: this runs per layer per
+        # predict, on a device that may be nearly full.
+        out = torch.empty(self.data.shape, dtype=torch.float16, device=self.data.device)
+        torch.mul(self.data, self.scale, out=out)
+        return out if self.dtype == torch.float16 else out.to(self.dtype)
 
     def slice(self, indices: Any) -> _Packed:
         # Scale has the same leading dims as data (last dim 1).
diff --git a/src/tabctx/engine.py b/src/tabctx/engine.py
index bc423b4..ce85062 100644
--- a/src/tabctx/engine.py
+++ b/src/tabctx/engine.py
@@ -12,7 +12,12 @@
 from tabctx.backends.base import TabularICLBackend
 from tabctx.cache.manager import CachedContext, ContextCacheManager
 from tabctx.chunking import choose_chunk_rows, split_rows
-from tabctx.errors import AdmissionRejected, DatasetNotFoundError, InvalidInputError
+from tabctx.errors import (
+    AdmissionRejected,
+    BackendComputeError,
+    DatasetNotFoundError,
+    InvalidInputError,
+)
 from tabctx.memory.estimator import MemoryEstimator
 from tabctx.memory.host import HostMemoryGuard
 from tabctx.types import ArrayLike, EngineStats, PredictOutcome, Task
@@ -83,6 +88,24 @@ def _validate_predict_input(X_test: ArrayLike, n_features: int) -> None:
         )
 
 
+def _predict_splitting_on_oom(
+    backend: TabularICLBackend, payload: object, X_test: ArrayLike, return_proba: bool
+) -> list[PredictOutcome]:
+    """Predict, halving the chunk on a CUDA OOM until it fits (or one row
+    still doesn't). Chunk sizing estimates can't see everything on a nearly
+    full device -- e.g. a compressed kv cache's per-layer decompression --
+    and a model may pick a single batch with no OOM retry of its own."""
+    try:
+        return [backend.predict(payload, X_test, return_proba=return_proba)]
+    except BackendComputeError:
+        if len(X_test) <= 1:
+            raise
+        half = len(X_test) // 2
+        return _predict_splitting_on_oom(
+            backend, payload, X_test[:half], return_proba
+        ) + _predict_splitting_on_oom(backend, payload, X_test[half:], return_proba)
+
+
 class TabctxEngine:
     """Orchestrates one or more backends over a single shared context
     cache and GPU budget. Multi-backend serving (e.g. TabICL and TabPFN
@@ -219,12 +242,7 @@ def fit(
         resolved_id = dataset_id or str(uuid.uuid4())
         # Host memory first, outside the cache lock: a large fit waiting for
         # a host slot must not block predicts on this device meanwhile.
-        host_slot = (
-            self._host_guard.reserve(n_train, n_features)
-            if self._host_guard is not None
-            else contextlib.nullcontext()
-        )
-        with host_slot, self._cache.lock:
+        with self._host_slot(chosen, n_train, n_features), self._cache.lock:
             # Usage-aware admission (v0.9.0): what OOMs is the fit's
             # transient PEAK on top of everything resident, so the
             # estimated peak must fit the headroom given current cache
@@ -277,9 +295,17 @@ def fit(
                 # only says how much was free), else the transient peak.
                 peak_hint = getattr(chosen, "fit_peak_bytes_hint", None)
                 peak_bytes = peak_hint() if callable(peak_hint) else None
+                # What the fit needed resident while it ran: backends that
+                # compress the context afterwards (TabICL's int8 kv cache)
+                # report the pre-compression size; est_bytes, the compressed
+                # size, stays the cache-accounting quantity.
+                fit_hint = getattr(chosen, "fit_resident_bytes_hint", None)
+                fit_resident = fit_hint() if callable(fit_hint) else None
                 record_fit = getattr(estimator, "record_fit", None)
                 if callable(record_fit):
-                    record_fit(n_train, n_features, est_bytes, peak_bytes)
+                    record_fit(
+                        n_train, n_features, fit_resident or est_bytes, peak_bytes
+                    )
                 else:
                     estimator.record_observation(
                         n_train, n_features, peak_bytes if peak_bytes else est_bytes
@@ -297,6 +323,20 @@ def fit(
             self._cache.put(context)
         return resolved_id
 
+    def _host_slot(self, backend: TabularICLBackend, n_rows: int, n_features: int):
+        """A process-wide large-call slot when this call may offload to host
+        memory (the backend's may_offload, else the guard's cell threshold);
+        a no-op without a guard."""
+        if self._host_guard is None:
+            return contextlib.nullcontext()
+        may_offload = getattr(backend, "may_offload", None)
+        large = (
+            may_offload(n_rows, n_features)
+            if callable(may_offload)
+            else self._host_guard.is_large(n_rows, n_features)
+        )
+        return self._host_guard.reserve(large)
+
     def feature_names(self, dataset_id: str) -> list[str] | None:
         """Column names recorded at fit() for a cached context, or None
         when the context is unknown or was fit without names. Reloads a
@@ -349,16 +389,20 @@ def predict(
         all_predictions: list = []
         all_probabilities: list[list[float]] = []
         classes: list[str] | None = None
-        with self._cache.lock:
+        largest_chunk = max(len(c) for c in chunks)
+        with (
+            self._host_slot(backend, largest_chunk, context.n_features),
+            self._cache.lock,
+        ):
             self._cache.touch(dataset_id)
             for chunk in chunks:
-                outcome = backend.predict(
-                    context.payload, chunk, return_proba=return_proba
-                )
-                all_predictions.extend(outcome.predictions)
-                if return_proba:
-                    all_probabilities.extend(outcome.probabilities or [])
-                    classes = outcome.classes or classes
+                for outcome in _predict_splitting_on_oom(
+                    backend, context.payload, chunk, return_proba
+                ):
+                    all_predictions.extend(outcome.predictions)
+                    if return_proba:
+                        all_probabilities.extend(outcome.probabilities or [])
+                        classes = outcome.classes or classes
 
         return PredictOutcome(
             predictions=all_predictions,
diff --git a/src/tabctx/memory/adaptive.py b/src/tabctx/memory/adaptive.py
index 30adf3e..7dd4452 100644
--- a/src/tabctx/memory/adaptive.py
+++ b/src/tabctx/memory/adaptive.py
@@ -140,10 +140,13 @@ def record_fit(
         resident_bytes: int,
         peak_bytes: int | None,
     ) -> None:
-        """Learn from a completed fit. With a working-floor model the
-        requirement is resident + floor (the peak of an adaptive model only
-        says how much was free); without one it is the peak, falling back
-        to the resident size for backends that don't report peaks."""
+        """Learn from a completed fit. resident_bytes is what the fit needed
+        resident while it ran (before any post-fit compression). With a
+        working-floor model the requirement is resident + floor -- not
+        capped at the measured peak, which under memory pressure only says
+        how little was free, while the floor is a measured minimum; without
+        one it is the peak, falling back to the resident size for backends
+        that don't report peaks."""
         if self._resident_rate is not None:
             self._resident_rate.record(n_train, n_features, resident_bytes)
         floor = (
@@ -153,8 +156,6 @@ def record_fit(
         )
         if floor is not None:
             need = resident_bytes + floor
-            if peak_bytes:
-                need = min(need, peak_bytes)  # never more than it actually used
         else:
             need = peak_bytes or resident_bytes
         self.record_observation(n_train, n_features, need)
diff --git a/src/tabctx/memory/calibration_tabicl_a100_requirement.py b/src/tabctx/memory/calibration_tabicl_a100_requirement.py
index 5426b00..eb54cd0 100644
--- a/src/tabctx/memory/calibration_tabicl_a100_requirement.py
+++ b/src/tabctx/memory/calibration_tabicl_a100_requirement.py
@@ -1,6 +1,6 @@
 """GENERATED by benchmarks/gen_requirement_grid.py -- do not hand-edit.
 
-Measured requirements (NVIDIA A100-SXM4-40GB, tabicl 2.2.0, kv mode, native KV precision, 2026-10-17): for each shape, the least free GPU
+Measured requirements (NVIDIA A100-SXM4-40GB, tabicl 2.2.0, kv mode, fit-time (native-precision) residents, 2026-10-17): for each shape, the least free GPU
 memory with which a fit succeeded (resident context + working floor), from
 benchmarks/bench_memory_pressure.py --search. See memory/requirement.py.
 """
@@ -12,7 +12,7 @@
 
 A100_40GB_TABICL_KV_FLOOR_POINTS: tuple[FloorPoint, ...] = (
     FloorPoint(
-        n_train=2000, n_features=10, working_bytes=214748364, resident_bytes=418759311
+        n_train=2000, n_features=10, working_bytes=322122547, resident_bytes=418759311
     ),  # resident 0.39 GiB, unconstrained peak 0.59 GiB, fit 0.433s -> 0.277s at the floor
     FloorPoint(
         n_train=20000,
@@ -23,13 +23,13 @@
     FloorPoint(
         n_train=45000,
         n_features=10,
-        working_bytes=4445291151,
+        working_bytes=8632884264,
         resident_bytes=7065221201,
     ),  # resident 6.58 GiB, unconstrained peak 10.72 GiB, fit 4.581s -> 4.051s at the floor
     FloorPoint(
         n_train=80000,
         n_features=10,
-        working_bytes=7902739824,
+        working_bytes=10157597655,
         resident_bytes=12058120683,
     ),  # resident 11.23 GiB, unconstrained peak 18.59 GiB, fit 10.753s -> 9.878s at the floor
     FloorPoint(
@@ -86,20 +86,20 @@
 )
 
 A100_40GB_TABICL_KV_REQUIREMENT_GRID: tuple[Observation, ...] = (
-    Observation(n_train=2000, n_features=10, real_bytes=633507675),
+    Observation(n_train=2000, n_features=10, real_bytes=740881858),
     Observation(n_train=20000, n_features=10, real_bytes=5712306503),
-    Observation(n_train=45000, n_features=10, real_bytes=11510512352),
-    Observation(n_train=80000, n_features=10, real_bytes=19960860507),
-    Observation(n_train=2000, n_features=50, real_bytes=1127428914),
+    Observation(n_train=45000, n_features=10, real_bytes=15698105466),
+    Observation(n_train=80000, n_features=10, real_bytes=22215718338),
+    Observation(n_train=2000, n_features=50, real_bytes=1127428915),
     Observation(n_train=10000, n_features=50, real_bytes=2985002270),
-    Observation(n_train=40000, n_features=50, real_bytes=12670153522),
+    Observation(n_train=40000, n_features=50, real_bytes=12670153523),
     Observation(n_train=2000, n_features=100, real_bytes=1202590842),
     Observation(n_train=20000, n_features=100, real_bytes=5851892940),
     Observation(n_train=50000, n_features=100, real_bytes=14001593384),
     Observation(n_train=1000, n_features=300, real_bytes=1309965025),
-    Observation(n_train=5000, n_features=300, real_bytes=2415919103),
+    Observation(n_train=5000, n_features=300, real_bytes=2415919104),
     Observation(n_train=1000, n_features=500, real_bytes=1621350154),
-    Observation(n_train=5000, n_features=500, real_bytes=3049426779),
-    Observation(n_train=500, n_features=1000, real_bytes=2609192631),
+    Observation(n_train=5000, n_features=500, real_bytes=3049426780),
+    Observation(n_train=500, n_features=1000, real_bytes=2609192632),
     Observation(n_train=2000, n_features=1000, real_bytes=3210488053),
 )
diff --git a/src/tabctx/memory/host.py b/src/tabctx/memory/host.py
index e54b5b1..fbaeb25 100644
--- a/src/tabctx/memory/host.py
+++ b/src/tabctx/memory/host.py
@@ -9,42 +9,78 @@
 at once, on two GPUs of one replica, get it OOM-killed with every tenant's
 context.
 
-HostMemoryGuard lets at most budget // large_fit_bytes large fits run at
-once, process-wide (shared by every device's engine). A large fit waits
-for a slot instead of exhausting the pod; small fits never wait. "Large"
-is a cell count below which no offload was ever measured.
+HostMemoryGuard lets at most budget // large_fit_bytes host-offloading
+calls run at once, process-wide (shared by every device's engine). Which
+calls may offload is the backend's call (TabICLBackend.may_offload mirrors
+tabicl's trigger: a stage output above half the device's currently free
+GPU memory); backends without that hint use a cell-count threshold. A
+large call waits for a slot instead of exhausting the pod; small ones
+never wait. Each slot's call should budget at most 1/slots of the host
+headroom for itself (see TabICLBackend's cpu_safety_factor), so
+concurrent calls can't all claim the same free memory.
 """
 
 from __future__ import annotations
 
 import contextlib
+import logging
 import threading
 from collections.abc import Iterator
 from pathlib import Path
 
 GiB = 1024**3
+logger = logging.getLogger(__name__)
 # Calibrated on an A100-40GB (see module docstring).
 LARGE_FIT_CELLS = 3_000_000
 LARGE_FIT_HOST_BYTES = 48 * GiB
 
 
-def detect_host_budget_bytes(fraction: float = 0.85) -> int | None:
-    """This process's memory limit (cgroup v2/v1, else physical RAM) times
-    fraction, or None when it can't be read."""
-    for path in (
-        "/sys/fs/cgroup/memory.max",
-        "/sys/fs/cgroup/memory/memory.limit_in_bytes",
-    ):
+CGROUP_LIMIT_PATHS = (
+    "/sys/fs/cgroup/memory.max",
+    "/sys/fs/cgroup/memory/memory.limit_in_bytes",
+)
+
+
+def detect_host_budget_bytes(
+    fraction: float = 0.85,
+    reserve_bytes: int = 8 * GiB,
+    limit_paths: tuple[str, ...] = CGROUP_LIMIT_PATHS,
+    meminfo_path: str = "/proc/meminfo",
+    status_path: str = "/proc/self/status",
+) -> int | None:
+    """Host memory available to large calls: this process's memory limit
+    (cgroup v2/v1, else physical RAM) x fraction, minus what the process
+    already holds (CUDA contexts, model weights) and a reserve for the
+    rest of serving (uploads, contexts' host state). None when the limit
+    can't be read."""
+    limit = _read_limit(limit_paths, meminfo_path)
+    if limit is None:
+        return None
+    return max(0, int(limit * fraction) - _rss_bytes(status_path) - reserve_bytes)
+
+
+def _rss_bytes(status_path: str) -> int:
+    try:
+        for line in Path(status_path).read_text().splitlines():
+            if line.startswith("VmRSS:"):
+                return int(line.split()[1]) * 1024
+    except OSError:
+        pass
+    return 0
+
+
+def _read_limit(limit_paths: tuple[str, ...], meminfo_path: str) -> int | None:
+    for path in limit_paths:
         try:
             raw = Path(path).read_text().strip()
         except OSError:
             continue
         if raw.isdigit() and int(raw) < 1 << 60:
-            return int(int(raw) * fraction)
+            return int(raw)
     try:
-        for line in Path("/proc/meminfo").read_text().splitlines():
+        for line in Path(meminfo_path).read_text().splitlines():
             if line.startswith("MemTotal:"):
-                return int(int(line.split()[1]) * 1024 * fraction)
+                return int(line.split()[1]) * 1024
     except OSError:
         pass
     return None
@@ -58,19 +94,29 @@ def __init__(
         large_fit_bytes: int = LARGE_FIT_HOST_BYTES,
     ) -> None:
         self.slots = max(1, budget_bytes // large_fit_bytes)
+        if budget_bytes < large_fit_bytes:
+            logger.warning(
+                "host memory budget %.1f GiB is below one large fit's measured "
+                "peak (%.0f GiB): large tables may exhaust this pod; give it more "
+                "memory or rely on disk offload (TABCTX_OFFLOAD_DIR)",
+                budget_bytes / GiB,
+                large_fit_bytes / GiB,
+            )
         self._large_fit_cells = large_fit_cells
         self._semaphore = threading.BoundedSemaphore(self.slots)
         self._waiting = 0
         self._active = 0
         self._lock = threading.Lock()
 
-    def is_large(self, n_train: int, n_features: int) -> bool:
-        return n_train * n_features >= self._large_fit_cells
+    def is_large(self, n_rows: int, n_features: int) -> bool:
+        """The fallback test for backends that can't say whether a call may
+        offload: a cell count below which no offload was ever measured."""
+        return n_rows * n_features >= self._large_fit_cells
 
     @contextlib.contextmanager
-    def reserve(self, n_train: int, n_features: int) -> Iterator[None]:
-        """Hold a large-fit slot for the duration (no-op for small fits)."""
-        if not self.is_large(n_train, n_features):
+    def reserve(self, large: bool) -> Iterator[None]:
+        """Hold a large-call slot for the duration (no-op when not large)."""
+        if not large:
             yield
             return
         with self._lock:
diff --git a/src/tabctx/memory/requirement.py b/src/tabctx/memory/requirement.py
index 416ebf2..7fcfaa1 100644
--- a/src/tabctx/memory/requirement.py
+++ b/src/tabctx/memory/requirement.py
@@ -27,7 +27,9 @@
 key/value per training row (plus a column cache independent of rows), so
 it is linear in rows: resident(n, d) <= n * max rate observed at >= d
 features. Rates are per-row bytes from real fits, which take precedence
-over calibration points (calibration may be at another precision).
+over calibration points. Residents are fit-time (pre-compression) sizes:
+a compressed context is smaller at rest, but the fit builds it at full
+precision first.
 """
 
 from __future__ import annotations
@@ -88,7 +90,14 @@ def bytes(self, n_train: int, n_features: int) -> int | None:
             return math.ceil(min(p.working_bytes for p in dominating) * self._margin)
         if n_features > self._max_features:
             return None
-        return math.ceil(self._base + self._per_row * n_train)
+        # Monotone in rows: never below a measured floor at >= these
+        # features and fewer rows (the per-row line ignores features).
+        shorter = [
+            p.working_bytes * self._margin
+            for p in self._points
+            if p.n_features >= n_features and p.n_train <= n_train
+        ]
+        return math.ceil(max([self._base + self._per_row * n_train, *shorter]))
 
 
 class ResidentRateModel:
@@ -100,9 +109,10 @@ class ResidentRateModel:
     features, resident r0), resident(n, d) <= r0 + rate * (n - n0) for
     n >= n0, where rate is the steepest per-row slope among long fits
     (>= slope_min_rows; short fits are dominated by the column cache).
-    The estimate is the smallest such bound. This replica's own fits feed
-    both the anchors and the slope, and supersede calibration slopes once
-    any long fit exists: calibration may be at another KV precision."""
+    The estimate is the smallest such bound. Residents are fit-time sizes
+    (before any post-fit compression -- what the fit needs while it runs),
+    the same quantity the calibration measured. This replica's own long
+    fits supersede the calibration slope once any exist."""
 
     def __init__(
         self, calibration: Iterable[FloorPoint] = (), slope_min_rows: int = 10_000
diff --git a/src/tabctx/serve/factory.py b/src/tabctx/serve/factory.py
index bbadbc9..1c1f6c2 100644
--- a/src/tabctx/serve/factory.py
+++ b/src/tabctx/serve/factory.py
@@ -412,7 +412,7 @@ def build_estimator(
 
 
 def _build_real_backend(
-    kind: BackendKind, settings: ServeSettings, device: str
+    kind: BackendKind, settings: ServeSettings, device: str, host_slots: int = 1
 ) -> TabularICLBackend:  # pragma: no cover
     # Requires torch + tabicl/tabpfn (GPU or licensed weights), which --
     # like backends/tabicl.py and backends/tabpfn.py themselves -- are not
@@ -438,6 +438,9 @@ def _build_real_backend(
             kv_cache=kv_cache,
             kv_precision=settings.kv_precision,
             offload_dir=_offload_dir(settings, device),
+            # Concurrent offloading calls each get their share of host
+            # headroom (see memory/host.py).
+            cpu_safety_factor=0.6 / host_slots,
         )
     # Weights load at replica start, not inside (and charged to) the first
     # tenant's fit.
@@ -459,7 +462,7 @@ def _offload_dir(settings: ServeSettings, device: str) -> str:
 
 
 def _build_backend(
-    kind: BackendKind, settings: ServeSettings, device: str = "cpu"
+    kind: BackendKind, settings: ServeSettings, device: str = "cpu", host_slots: int = 1
 ) -> TabularICLBackend:
     """One backend bound to `device`. Imports torch/tabicl/tabpfn only on
     the paths that need them, so the fake backend runs with core deps
@@ -468,7 +471,7 @@ def _build_backend(
         from tabctx.backends.fake import FakeBackend
 
         return FakeBackend()
-    return _build_real_backend(kind, settings, device)
+    return _build_real_backend(kind, settings, device, host_slots)
 
 
 def _warn_on_spill_layout(spill_dir: str, multi: bool) -> None:
@@ -541,7 +544,9 @@ def build_engine(settings: ServeSettings | None = None) -> BuiltEngine:
     for device in devices:
         backends: dict[str, TabularICLBackend] = {}
         for kind in settings.backends:
-            backend = _build_backend(kind, settings, device)
+            backend = _build_backend(
+                kind, settings, device, host_guard.slots if host_guard else 1
+            )
             backends[backend.name] = backend
             # One estimator per model, shared by every device: they are
             # identical cards budgeted alike, and observations learned on
diff --git a/tests/integration/test_tabicl_kv_real.py b/tests/integration/test_tabicl_kv_real.py
new file mode 100644
index 0000000..0770cbb
--- /dev/null
+++ b/tests/integration/test_tabicl_kv_real.py
@@ -0,0 +1,55 @@
+"""CompressedKVEntry against real tabicl structures (CPU): fit, compress,
+predict, pickle round trip -- the paths a hand-written stand-in can't
+prove (KVCache batching, TabICLCache.to, __getstate__/__setstate__).
+
+Skipped without tabicl or its checkpoint in the local HuggingFace cache."""
+
+import pickle
+
+import numpy as np
+import pytest
+
+pytest.importorskip("tabicl")
+torch = pytest.importorskip("torch")
+
+from tabctx.backends.tabicl import TabICLBackend  # noqa: E402
+from tabctx.backends.tabicl_kv import CompressedKVEntry  # noqa: E402
+
+
+@pytest.fixture(scope="module")
+def data():
+    rng = np.random.default_rng(0)
+    X = rng.normal(size=(1200, 6)).astype(np.float32)  # >1024 rows: AMP path
+    return X, (X[:, 0] + X[:, 1] > 0).astype(int)
+
+
+def _backend(precision):
+    try:
+        b = TabICLBackend(device="cpu", kv_precision=precision)
+        b.warmup()
+    except Exception as e:  # noqa: BLE001 -- no checkpoint offline
+        pytest.skip(f"TabICL checkpoint unavailable: {e}")
+    return b
+
+
+@pytest.mark.parametrize("precision", ["int8", "fp16"])
+def test_compressed_context_predicts_like_native(data, precision):
+    X, y = data
+    native, comp = _backend("native"), _backend(precision)
+    want = native.predict(
+        native.fit(X[:1000], y[:1000], "classification"), X[1000:], True
+    )
+    payload = comp.fit(X[:1000], y[:1000], "classification")
+    entries = [
+        e for c in payload.model_kv_cache_.values() for e in c.icl_cache.kv.values()
+    ]
+    assert entries and all(isinstance(e, CompressedKVEntry) for e in entries)
+    got = comp.predict(payload, X[1000:], True)
+    assert got.predictions == want.predictions
+    assert (
+        np.abs(np.array(got.probabilities) - np.array(want.probabilities)).max() < 0.02
+    )
+    # Spill round trip through tabicl's own pickling (cache.to("cpu") etc.).
+    restored = comp.loads_payload(comp.dumps_payload(payload))
+    assert comp.predict(restored, X[1000:], True).predictions == got.predictions
+    assert pickle.loads(pickle.dumps(entries[0])).is_valid()
diff --git a/tests/unit/test_engine.py b/tests/unit/test_engine.py
index d988b17..79472e5 100644
--- a/tests/unit/test_engine.py
+++ b/tests/unit/test_engine.py
@@ -273,3 +273,72 @@ def test_default_backend_property(self):
     def test_estimator_for_default_and_named(self):
         engine, _ = make_engine()
         assert engine.estimator_for() is engine.estimator_for("fake")
+
+
+class TestPredictOomFallback:
+    def test_an_oom_chunk_is_split_until_it_fits(self):
+        from tabctx.backends.fake import FakeBackend
+        from tabctx.cache.manager import ContextCacheManager
+        from tabctx.engine import TabctxEngine
+        from tabctx.errors import BackendComputeError
+        from tabctx.memory import (
+            A100_40GB_TABICL_CALIBRATION,
+            AdaptiveMemoryEstimator,
+            PowerLawMemoryEstimator,
+        )
+
+        class OomAbove(FakeBackend):
+            def __init__(self, max_rows):
+                super().__init__(bytes_hint=10)
+                self.max_rows = max_rows
+                self.calls = []
+
+            def predict(self, payload, X_test, return_proba=False):
+                self.calls.append(len(X_test))
+                if len(X_test) > self.max_rows:
+                    raise BackendComputeError("CUDA OOM during predict()")
+                return super().predict(payload, X_test, return_proba)
+
+        backend = OomAbove(max_rows=3)
+        est = AdaptiveMemoryEstimator(
+            fallback=PowerLawMemoryEstimator(A100_40GB_TABICL_CALIBRATION)
+        )
+        engine = TabctxEngine(
+            backend=backend,
+            cache=ContextCacheManager(est.ceiling_bytes()),
+            estimator=est,
+        )
+        engine.fit([[1.0], [2.0]], ["a", "a"], dataset_id="d")
+        out = engine.predict("d", [[0.0]] * 10)
+        assert out.predictions == ["a"] * 10  # all rows, in order
+        served = [n for n in backend.calls if n <= 3]  # the calls that fit
+        assert sum(served) == 10 and backend.calls[0] == 10
+
+    def test_a_single_row_that_still_ooms_raises(self):
+        import pytest
+
+        from tabctx.backends.fake import FakeBackend
+        from tabctx.cache.manager import ContextCacheManager
+        from tabctx.engine import TabctxEngine
+        from tabctx.errors import BackendComputeError
+        from tabctx.memory import (
+            A100_40GB_TABICL_CALIBRATION,
+            AdaptiveMemoryEstimator,
+            PowerLawMemoryEstimator,
+        )
+
+        class AlwaysOom(FakeBackend):
+            def predict(self, payload, X_test, return_proba=False):
+                raise BackendComputeError("CUDA OOM during predict()")
+
+        est = AdaptiveMemoryEstimator(
+            fallback=PowerLawMemoryEstimator(A100_40GB_TABICL_CALIBRATION)
+        )
+        engine = TabctxEngine(
+            backend=AlwaysOom(bytes_hint=10),
+            cache=ContextCacheManager(est.ceiling_bytes()),
+            estimator=est,
+        )
+        engine.fit([[1.0]], ["a"], dataset_id="d")
+        with pytest.raises(BackendComputeError):
+            engine.predict("d", [[0.0]] * 4)
diff --git a/tests/unit/test_host_memory.py b/tests/unit/test_host_memory.py
index c056632..cc4fd11 100644
--- a/tests/unit/test_host_memory.py
+++ b/tests/unit/test_host_memory.py
@@ -20,7 +20,7 @@ def test_slots_from_budget():
 
 def test_small_fits_never_wait():
     guard = HostMemoryGuard(GiB, large_fit_cells=100, large_fit_bytes=GiB)
-    with guard.reserve(50, 1), guard.reserve(50, 1):  # both small: no slot used
+    with guard.reserve(False), guard.reserve(False):  # neither large: no slot used
         assert guard.stats()["large_fits_active"] == 0
 
 
@@ -31,7 +31,7 @@ def test_large_fits_are_limited_to_the_slots():
     inside, release = threading.Semaphore(0), threading.Event()
 
     def large():
-        with guard.reserve(100, 1):
+        with guard.reserve(True):
             inside.release()
             release.wait(5)
 
@@ -97,3 +97,101 @@ def test_a_waiting_large_fit_does_not_block_predicts_on_its_device():
     backends[1].gate.set()
     t0.join(5)
     t1.join(5)
+
+
+def test_budget_from_cgroup_v2_minus_baseline_rss(tmp_path):
+    from tabctx.memory.host import detect_host_budget_bytes
+
+    (tmp_path / "memory.max").write_text(f"{64 * GiB}\n")
+    (tmp_path / "status").write_text("Name: python\nVmRSS:\t 4194304 kB\n")  # 4 GiB
+    budget = detect_host_budget_bytes(
+        fraction=0.85,
+        reserve_bytes=8 * GiB,
+        limit_paths=(str(tmp_path / "memory.max"),),
+        meminfo_path=str(tmp_path / "none"),
+        status_path=str(tmp_path / "status"),
+    )
+    assert budget == int(64 * GiB * 0.85) - 4 * GiB - 8 * GiB
+
+
+def test_unlimited_cgroup_falls_back_to_physical_memory(tmp_path):
+    from tabctx.memory.host import detect_host_budget_bytes
+
+    (tmp_path / "v2").write_text("max\n")
+    (tmp_path / "v1").write_text("9223372036854771712\n")  # v1 "unlimited"
+    (tmp_path / "meminfo").write_text("MemTotal:       16777216 kB\n")  # 16 GiB
+    budget = detect_host_budget_bytes(
+        fraction=1.0,
+        reserve_bytes=0,
+        limit_paths=(str(tmp_path / "v2"), str(tmp_path / "v1")),
+        meminfo_path=str(tmp_path / "meminfo"),
+        status_path=str(tmp_path / "none"),
+    )
+    assert budget == 16 * GiB
+
+
+def test_unreadable_limit_means_no_budget(tmp_path):
+    from tabctx.memory.host import detect_host_budget_bytes
+
+    assert (
+        detect_host_budget_bytes(
+            limit_paths=(str(tmp_path / "x"),), meminfo_path=str(tmp_path / "y")
+        )
+        is None
+    )
+
+
+def test_budget_below_one_large_fit_warns(caplog):
+    with caplog.at_level("WARNING"):
+        assert HostMemoryGuard(20 * GiB, large_fit_bytes=48 * GiB).slots == 1
+    assert "below one large fit" in caplog.text
+
+
+class _Offloading(FakeBackend):
+    """A backend that says its calls offload (like TabICL on a busy card)."""
+
+    def __init__(self, offloads: bool) -> None:
+        super().__init__(bytes_hint=10)
+        self.offloads = offloads
+        self.queries: list[tuple[int, int]] = []
+
+    def may_offload(self, n_rows, n_features):
+        self.queries.append((n_rows, n_features))
+        return self.offloads
+
+
+def _engine(backend, guard):
+    est = AdaptiveMemoryEstimator(
+        fallback=PowerLawMemoryEstimator(A100_40GB_TABICL_CALIBRATION)
+    )
+    return TabctxEngine(
+        backend=backend,
+        cache=ContextCacheManager(est.ceiling_bytes()),
+        estimator=est,
+        host_guard=guard,
+    )
+
+
+def test_backend_hint_decides_and_predicts_are_guarded_too():
+    guard = HostMemoryGuard(GiB, large_fit_cells=10**12, large_fit_bytes=GiB)
+    backend = _Offloading(offloads=True)
+    engine = _engine(backend, guard)
+    seen = []
+    real = guard.reserve
+    guard.reserve = lambda large: (seen.append(large), real(large))[1]
+    engine.fit([[1.0, 2.0]] * 3, ["a"] * 3, dataset_id="d")
+    engine.predict("d", [[1.0, 2.0]] * 5)
+    # The hint, not the (huge) cell threshold, decided -- for fit AND predict.
+    assert seen == [True, True]
+    assert backend.queries == [(3, 2), (5, 2)]
+
+
+def test_without_a_hint_the_cell_threshold_decides():
+    guard = HostMemoryGuard(GiB, large_fit_cells=6, large_fit_bytes=GiB)
+    engine = _engine(FakeBackend(bytes_hint=10), guard)
+    seen = []
+    real = guard.reserve
+    guard.reserve = lambda large: (seen.append(large), real(large))[1]
+    engine.fit([[1.0, 2.0]] * 3, ["a"] * 3, dataset_id="d")  # 6 cells: large
+    engine.predict("d", [[1.0, 2.0]])  # 2 cells: not
+    assert seen == [True, False]
diff --git a/tests/unit/test_requirement.py b/tests/unit/test_requirement.py
index debb5a0..90885e9 100644
--- a/tests/unit/test_requirement.py
+++ b/tests/unit/test_requirement.py
@@ -1,6 +1,8 @@
 """Unit tests for requirement-based admission (memory/requirement.py):
 fits are learned as resident + working floor, not unconstrained peaks."""
 
+import math
+
 import pytest
 
 from tabctx.backends.fake import FakeBackend
@@ -67,11 +69,13 @@ def test_with_a_floor_model_learns_resident_plus_floor(self):
         expected = (4 + 1.6) * GiB * 1.5  # runtime safety margin
         assert est.estimate_bytes(20_000, 0, 100) == pytest.approx(expected, rel=1e-3)
 
-    def test_never_more_than_the_measured_peak(self):
+    def test_not_capped_at_a_peak_measured_under_pressure(self):
+        # A peak below resident + floor only says how little was free; the
+        # floor is a measured minimum, so the requirement keeps it.
         est = _estimator(WorkingFloorModel(POINTS, margin=1.0))
         est.record_fit(20_000, 100, resident_bytes=4 * GiB, peak_bytes=int(4.5 * GiB))
         assert est.estimate_bytes(20_000, 0, 100) == pytest.approx(
-            4.5 * GiB * 1.5, rel=1e-3
+            5.6 * GiB * 1.5, rel=1e-3
         )
 
     def test_without_a_floor_model_learns_the_peak(self):
@@ -95,18 +99,15 @@ def __init__(self, resident: int, peak: int) -> None:
 
 class TestAdmissionEffect:
     def test_wide_table_is_admitted_without_evicting_other_tenants(self):
-        """The motivating case: a 20k x 100 fit's unconstrained peak is 25 GB,
-        but it needs only resident + floor (5.6 GB). Each estimator is
-        calibrated as the factory does (a measured requirement vs a measured
-        peak). Beside a warm cache of five 4 GB tenants, peak semantics evict
-        tenants to make room the fit doesn't need; requirement semantics
-        admit it untouched."""
+        """The motivating case. Both estimators start from the same measured
+        PEAK (25 GB for 20k x 100, as the real peak grid has); the only
+        difference is the working-floor model. With it, the first fit is
+        learned as resident + floor (4 + 2 GB), so the second fit beside five
+        4 GB tenants needs ~9 GB and evicts nobody; without it, the learned
+        peak makes the fit evict tenants for room it doesn't need."""
 
         def run(floor):
-            measured = (5.6 if floor else 25) * GiB
-            est = _estimator(
-                floor, preloaded=(Observation(20_000, 100, int(measured)),)
-            )
+            est = _estimator(floor, preloaded=(Observation(20_000, 100, 25 * GiB),))
             backend = _AdaptiveFake(resident=4 * GiB, peak=25 * GiB)
             cache = ContextCacheManager(capacity_bytes=34 * GiB)
             engine = TabctxEngine(backend=backend, cache=cache, estimator=est)
@@ -124,3 +125,104 @@ def run(floor):
 
         assert len(run(WorkingFloorModel(POINTS))) == 5  # nobody evicted
         assert len(run(None)) < 5  # peak semantics evict tenants
+
+
+class TestResidentRateModel:
+    CAL = (
+        FloorPoint(20_000, 10, GiB, resident_bytes=4 * GiB),
+        FloorPoint(80_000, 10, GiB, resident_bytes=16 * GiB),
+        FloorPoint(2_000, 1_000, GiB, resident_bytes=2 * GiB),
+    )
+
+    def test_extrapolates_linearly_from_the_nearest_anchor(self):
+        from tabctx.memory.requirement import ResidentRateModel
+
+        m = ResidentRateModel(self.CAL)
+        slope = 16 * GiB / 80_000  # steepest long fit
+        assert m.estimate(100_000, 10) == pytest.approx(
+            16 * GiB + slope * 20_000, rel=1e-6
+        )
+
+    def test_short_wide_fits_do_not_set_the_slope(self):
+        from tabctx.memory.requirement import ResidentRateModel
+
+        m = ResidentRateModel(self.CAL)
+        # 2000x1000 has 1 MiB/row (column cache) but is too short to set it.
+        assert m.estimate(100_000, 10) < 25 * GiB
+
+    def test_wider_query_uses_only_anchors_at_least_as_wide(self):
+        from tabctx.memory.requirement import ResidentRateModel
+
+        m = ResidentRateModel(self.CAL)
+        slope = 16 * GiB / 80_000
+        assert m.estimate(10_000, 500) == pytest.approx(
+            2 * GiB + slope * 8_000, rel=1e-6
+        )
+        assert m.estimate(10, 5_000) is None
+
+    def test_runtime_long_fits_supersede_the_calibration_slope(self):
+        from tabctx.memory.requirement import ResidentRateModel
+
+        m = ResidentRateModel(self.CAL)
+        m.record(50_000, 10, 5 * GiB)  # a replica fit: 0.1 MiB/row
+        assert m.estimate(100_000, 10) == pytest.approx(5 * GiB + 5 * GiB, rel=1e-6)
+
+    def test_runtime_history_is_bounded(self):
+        from tabctx.memory.requirement import ResidentRateModel
+
+        m = ResidentRateModel()
+        for i in range(700):
+            m.record(10_000 + i, 10, GiB)
+        assert len(m._runtime) == 500
+
+
+class TestExtrapolatedRequirement:
+    def test_beyond_the_grid_uses_resident_plus_floor_not_the_formula(self):
+        from tabctx.memory.requirement import ResidentRateModel
+
+        floor = WorkingFloorModel(POINTS)
+        rate = ResidentRateModel(TestResidentRateModel.CAL)
+        est = AdaptiveMemoryEstimator(
+            fallback=PowerLawMemoryEstimator(A100_40GB_TABICL_CALIBRATION),
+            working_floor=floor,
+            resident_rate=rate,
+            resident_margin=1.2,
+        )
+        n, d = 120_000, 10
+        expected = math.ceil(rate.estimate(n, d) * 1.2) + floor.bytes(n, d)
+        assert est.estimate_bytes(n, 0, d) == expected
+
+    def test_wider_than_any_measurement_falls_back(self):
+        from tabctx.memory.requirement import ResidentRateModel
+
+        est = AdaptiveMemoryEstimator(
+            fallback=PowerLawMemoryEstimator(A100_40GB_TABICL_CALIBRATION),
+            working_floor=WorkingFloorModel(POINTS),
+            resident_rate=ResidentRateModel(TestResidentRateModel.CAL),
+        )
+        assert est.estimate_bytes(100, 0, 5_000) == est._fallback.estimate_bytes(
+            100, 0, 5_000
+        )
+
+
+def test_floor_model_is_monotone_in_rows():
+    m = WorkingFloorModel(
+        (
+            FloorPoint(2_000, 1_000, int(1.15 * GiB)),
+            FloorPoint(20_000, 100, int(1.6 * GiB)),
+        )
+    )
+    assert m.bytes(3_000, 1_000) >= m.bytes(2_000, 1_000)
+
+
+def test_factory_wires_floor_and_resident_models_for_tabicl_kv_only():
+    from tabctx.serve.factory import ServeSettings, build_estimator
+
+    kv = build_estimator(ServeSettings(backends=("tabicl",)), "tabicl", 40 * GiB)
+    assert kv._working_floor is not None and kv._resident_rate is not None
+    for settings, kind in (
+        (ServeSettings(backends=("tabicl",), kv_cache="repr"), "tabicl"),
+        (ServeSettings(backends=("tabpfn",)), "tabpfn"),
+    ):
+        est = build_estimator(settings, kind, 40 * GiB)
+        assert est._working_floor is None and est._resident_rate is None
diff --git a/tests/unit/test_serve_factory.py b/tests/unit/test_serve_factory.py
index 7f8372f..c6473eb 100644
--- a/tests/unit/test_serve_factory.py
+++ b/tests/unit/test_serve_factory.py
@@ -170,7 +170,9 @@ def test_none_capacity_keeps_the_a100_defaults(self):
     def test_build_engine_uses_detected_capacity_for_real_backends(self, monkeypatch):
         monkeypatch.setattr(
             "tabctx.serve.factory._build_real_backend",
-            lambda kind, settings, device: FakeBackend(name=f"{kind}-model"),
+            lambda kind, settings, device, host_slots=1: FakeBackend(
+                name=f"{kind}-model"
+            ),
         )
         monkeypatch.setattr(
             "tabctx.serve.factory.detect_gpu_capacity_bytes",
@@ -232,7 +234,9 @@ def test_default_model_is_the_first_listed_kind_by_model_id(self, monkeypatch):
         # model, resolved through that mapping.
         monkeypatch.setattr(
             "tabctx.serve.factory._build_real_backend",
-            lambda kind, settings, device: FakeBackend(name=f"{kind}-model"),
+            lambda kind, settings, device, host_slots=1: FakeBackend(
+                name=f"{kind}-model"
+            ),
         )
         built = build_engine(
             ServeSettings(backends=("tabpfn", "fake"), devices=("cpu",))
@@ -276,7 +280,7 @@ def test_dispatches_non_fake_kinds_to_the_real_backend_path(self, monkeypatch):
         calls = []
         monkeypatch.setattr(
             "tabctx.serve.factory._build_real_backend",
-            lambda kind, settings, device: (
+            lambda kind, settings, device, host_slots=1: (
                 calls.append((kind, settings, device)) or object()
             ),
         )
@@ -383,7 +387,7 @@ def test_several_devices_build_one_engine_each(self, monkeypatch, tmp_path):
         created = []
         monkeypatch.setattr(
             "tabctx.serve.factory._build_real_backend",
-            lambda kind, settings, device: (
+            lambda kind, settings, device, host_slots=1: (
                 created.append(device) or FakeBackend(name=f"{kind}-model")
             ),
         )
diff --git a/tests/unit/test_tabicl_host_release.py b/tests/unit/test_tabicl_host_release.py
new file mode 100644
index 0000000..ab5f041
--- /dev/null
+++ b/tests/unit/test_tabicl_host_release.py
@@ -0,0 +1,70 @@
+"""TabICLBackend releases offload staging memory -- also when a fit fails
+(the review's finding: a failed large fit kept its pinned buffers while its
+host slot went to the next fit)."""
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+from tabctx.backends.tabicl import TabICLBackend  # noqa: E402
+
+
+class _Pool:
+    def __init__(self):
+        self.cleared = 0
+
+    def clear(self):
+        self.cleared += 1
+
+
+class _Stage:
+    def __init__(self):
+        self.inference_mgr = type("Mgr", (), {"_buffer_pool": _Pool()})()
+
+
+class _Module:
+    def __init__(self):
+        self.stages = [_Stage(), _Stage()]
+
+    def modules(self):
+        return iter([self, *self.stages])
+
+    def clear_cache(self):
+        pass
+
+
+class _Estimator:
+    def __init__(self, fail: bool):
+        self.model_ = _Module()
+        self.fail = fail
+
+    def fit(self, X, y):
+        if self.fail:
+            raise RuntimeError("Failed to execute even with minimum batch size 1")
+
+
+def _backend(estimator):
+    b = object.__new__(TabICLBackend)  # no tabicl needed
+    b._device = "cpu"
+    b._kv_precision = "native"
+    b._shared_backbones = {}
+    b._make_estimator = lambda task: estimator
+    b._stash_backbone = lambda task, model: None
+    return b
+
+
+def pools(est):
+    return [s.inference_mgr._buffer_pool.cleared for s in est.model_.stages]
+
+
+def test_pools_are_cleared_after_a_fit():
+    est = _Estimator(fail=False)
+    _backend(est).fit([[1.0]], [0], "classification")
+    assert pools(est) == [1, 1]
+
+
+def test_pools_are_cleared_even_when_the_fit_fails():
+    est = _Estimator(fail=True)
+    with pytest.raises(RuntimeError):
+        _backend(est).fit([[1.0]], [0], "classification")
+    assert pools(est) == [1, 1]

From cb037e35dccd129aaa8dd9327a845934c9b1e09e Mon Sep 17 00:00:00 2001
From: Amrit Krishnan <amrit110@gmail.com>
Date: Sat, 17 Oct 2026 17:03:03 -0400
Subject: [PATCH 08/39] fix(engine): count each device's baseline memory
 (weights, CUDA context) in admission

Measured after warmup as total - free. On a 4x L4 multi-GPU rehearsal,
fits the books said were safe OOMed: ~2 GB of weights + context per
device isn't in cache accounting and the 15% reserve of a 22 GB card
doesn't absorb it.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/tabctx/engine.py        | 26 ++++++++++++++++--------
 src/tabctx/serve/factory.py | 21 ++++++++++++++++++-
 tests/unit/test_engine.py   | 40 +++++++++++++++++++++++++++++++++++++
 3 files changed, 78 insertions(+), 9 deletions(-)

diff --git a/src/tabctx/engine.py b/src/tabctx/engine.py
index ce85062..3f710b6 100644
--- a/src/tabctx/engine.py
+++ b/src/tabctx/engine.py
@@ -125,6 +125,7 @@ def __init__(
         estimators: dict[str, MemoryEstimator] | None = None,
         default_backend: str | None = None,
         host_guard: HostMemoryGuard | None = None,
+        device_baseline_bytes: int = 0,
     ) -> None:
         """Single-backend form: TabctxEngine(backend=, cache=, estimator=).
         Multi-backend form: TabctxEngine(backends={name: b}, cache=,
@@ -151,6 +152,13 @@ def __init__(
         # Process-wide host-memory admission, shared by every device's
         # engine (see memory/host.py); None = unguarded.
         self._host_guard = host_guard
+        # Device memory in use before any context: CUDA context, model
+        # weights (TabPFN's shared models ~1.7 GB, TabICL's backbone),
+        # allocator overhead -- invisible to the cache's accounting.
+        # Measured after warmup; counted against every headroom check (on a
+        # 22 GB L4 the 15% transient reserve doesn't absorb it: fits the
+        # books said were safe OOMed).
+        self._baseline_bytes = device_baseline_bytes
 
     @property
     def backend_names(self) -> list[str]:
@@ -167,10 +175,13 @@ def admission_headroom_bytes(self, backend: str | None = None) -> int:
         """Transient headroom a fit could use right now, given what the
         cache holds resident (the quantity fit admission bounds).
         Lock-free snapshot: never waits on this engine's GPU work."""
-        return self.estimator_for(backend).admission_headroom_bytes(
-            self._cache.stats_nowait().used_bytes
+        return self._headroom(
+            self.estimator_for(backend), self._cache.stats_nowait().used_bytes
         )
 
+    def _headroom(self, estimator: MemoryEstimator, used_bytes: int) -> int:
+        return estimator.admission_headroom_bytes(used_bytes + self._baseline_bytes)
+
     def holds(self, dataset_id: str) -> bool:
         """Resident in this engine's cache (lock-free snapshot; a
         spilled copy doesn't count)."""
@@ -262,18 +273,17 @@ def fit(
             # anything -- otherwise one oversized request drains every
             # tenant's context on its way to a 413 (observed live: a
             # 200k x 60 request spilled 30 contexts and then failed).
-            if estimated > estimator.admission_headroom_bytes(0):
+            max_headroom = self._headroom(estimator, 0)
+            if estimated > max_headroom:
                 raise AdmissionRejected(
                     f"training shape ({n_train} rows x {n_features} "
                     f"features) is estimated to need {estimated} bytes at "
                     f"peak, exceeding this replica's "
-                    f"{estimator.admission_headroom_bytes(0)} byte maximum "
+                    f"{max_headroom} byte maximum "
                     "admission headroom (even with an empty cache)"
                 )
             while True:
-                headroom = estimator.admission_headroom_bytes(
-                    self._cache.stats().used_bytes
-                )
+                headroom = self._headroom(estimator, self._cache.stats().used_bytes)
                 if estimated <= headroom:
                     break
                 # Feasible but blocked by warm contexts: evict-ahead.
@@ -382,7 +392,7 @@ def predict(
             estimator,
             context.n_train,
             context.n_features,
-            estimator.admission_headroom_bytes(self._cache.stats().used_bytes),
+            self._headroom(estimator, self._cache.stats().used_bytes),
         )
         chunks = split_rows(X_test, chunk_rows) if chunk_rows < n_test else [X_test]
 
diff --git a/src/tabctx/serve/factory.py b/src/tabctx/serve/factory.py
index 1c1f6c2..7896d46 100644
--- a/src/tabctx/serve/factory.py
+++ b/src/tabctx/serve/factory.py
@@ -340,6 +340,22 @@ def _visible_cuda_count() -> int:
     return torch.cuda.device_count() if torch.cuda.is_available() else 0
 
 
+def measure_device_baseline(device: str) -> int:
+    """Device memory in use before any context (call after the backends'
+    warmup): total - free, so it includes the CUDA context and allocator
+    overhead as well as the model weights. 0 off CUDA."""
+    if not device.startswith("cuda"):
+        return 0
+    try:
+        import torch
+    except ImportError:
+        return 0
+    if not torch.cuda.is_available():
+        return 0
+    free, total = torch.cuda.mem_get_info(torch.device(device))
+    return int(total - free)
+
+
 def resolve_devices(settings: ServeSettings) -> tuple[str, ...]:
     """The devices to build one engine for each. Explicit settings win
     (GPU indices are checked against what this process can see); "auto"
@@ -530,6 +546,7 @@ def build_engine(settings: ServeSettings | None = None) -> BuiltEngine:
         else detect_gpu_capacity_bytes(devices)
     )
     multi = len(devices) > 1
+    fake_only = set(settings.backends) == {"fake"}
     if settings.spill_dir:
         _warn_on_spill_layout(settings.spill_dir, multi)
     estimators: dict[str, MemoryEstimator] = {}
@@ -575,6 +592,9 @@ def build_engine(settings: ServeSettings | None = None) -> BuiltEngine:
                 cache=cache,
                 default_backend=default,
                 host_guard=host_guard,
+                device_baseline_bytes=0
+                if fake_only
+                else measure_device_baseline(device),
             )
         )
         spill_stores.append(spill_store)
@@ -582,7 +602,6 @@ def build_engine(settings: ServeSettings | None = None) -> BuiltEngine:
     engine: TabctxEngine | MultiDeviceEngine = (
         MultiDeviceEngine(shards, list(devices)) if multi else shards[0]
     )
-    fake_only = set(settings.backends) == {"fake"}
     return BuiltEngine(
         engine=engine,
         estimators={name: estimators[name] for name in first_backends},
diff --git a/tests/unit/test_engine.py b/tests/unit/test_engine.py
index 79472e5..e57545c 100644
--- a/tests/unit/test_engine.py
+++ b/tests/unit/test_engine.py
@@ -342,3 +342,43 @@ def predict(self, payload, X_test, return_proba=False):
         engine.fit([[1.0]], ["a"], dataset_id="d")
         with pytest.raises(BackendComputeError):
             engine.predict("d", [[0.0]] * 4)
+
+
+def test_device_baseline_counts_against_headroom():
+    """Weights + CUDA context occupy the device before any context; a fit
+    that fits an empty card must be rejected when they leave too little."""
+    import pytest
+
+    from tabctx.backends.fake import FakeBackend
+    from tabctx.cache.manager import ContextCacheManager
+    from tabctx.engine import TabctxEngine
+    from tabctx.errors import AdmissionRejected
+    from tabctx.memory import (
+        A100_40GB_TABICL_CALIBRATION,
+        AdaptiveMemoryEstimator,
+        PowerLawMemoryEstimator,
+    )
+    from tabctx.memory.adaptive import Observation
+
+    GiB = 1024**3
+
+    def engine(baseline):
+        fallback = PowerLawMemoryEstimator(
+            A100_40GB_TABICL_CALIBRATION,
+            hard_ceiling_bytes=12 * GiB,
+            gpu_capacity_bytes=22 * GiB,
+        )
+        est = AdaptiveMemoryEstimator(
+            fallback=fallback, preloaded=(Observation(1_000, 2, int(15 * GiB / 1.1)),)
+        )
+        return TabctxEngine(
+            backend=FakeBackend(bytes_hint=GiB),
+            cache=ContextCacheManager(12 * GiB),
+            estimator=est,
+            device_baseline_bytes=baseline,
+        )
+
+    X, y = [[1.0, 2.0]] * 1_000, ["a"] * 1_000
+    engine(0).fit(X, y)  # 15 GB need <= 0.85 x 22 = 18.7 GB
+    with pytest.raises(AdmissionRejected):
+        engine(4 * GiB).fit(X, y)  # 18.7 - 4 = 14.7 GB < 15 GB

From fd2227e79ae350420d582043ef348ef648dad24f Mon Sep 17 00:00:00 2001
From: Amrit Krishnan <amrit110@gmail.com>
Date: Sat, 17 Oct 2026 18:03:55 -0400
Subject: [PATCH 09/39] fix(multidevice): contexts moving between tiers stay
 routable

A fit returned 200 and the next predict 404'd within 30 ms on the 4xL4
rehearsal: a concurrent fit's evict-ahead was spilling the new context,
and while the spill I/O ran it was listed by neither the GPU cache nor
the spill store, so the router's lock-free lookup found nothing. The
cache now tracks ids in transit (being spilled or restored) and has()
reports them; the shard's locked get() waits for the move.

predict also looks once more after missing everywhere: a re-fit can move
an id to another device between the router's snapshot and its calls.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/tabctx/cache/manager.py    | 54 +++++++++++++++++++----------
 src/tabctx/multidevice.py      | 26 +++++++++-----
 tests/unit/test_multidevice.py | 48 ++++++++++++++++++++++++++
 tests/unit/test_spill.py       | 62 ++++++++++++++++++++++++++++++++++
 4 files changed, 164 insertions(+), 26 deletions(-)

diff --git a/src/tabctx/cache/manager.py b/src/tabctx/cache/manager.py
index de50c59..0b34302 100644
--- a/src/tabctx/cache/manager.py
+++ b/src/tabctx/cache/manager.py
@@ -67,6 +67,11 @@ def __init__(
         self._used_bytes = 0
         self._spill = spill_store
         self._lock = threading.RLock()
+        # Ids moving between tiers (being spilled, or restored), which
+        # neither tier lists until the I/O finishes. Lock-free readers
+        # (has) must still see them: a multi-device router that saw
+        # neither would answer 404 for a context that is merely moving.
+        self._in_transit: set[str] = set()
 
     def get(self, dataset_id: str) -> CachedContext | None:
         with self._lock:
@@ -76,19 +81,23 @@ def get(self, dataset_id: str) -> CachedContext | None:
             # Miss with a spill tier: try reloading. Re-admitting evicts
             # via the normal policy to make room (RLock, so the nested
             # put() is fine); what it displaces may itself spill.
-            restored = self._spill.load(dataset_id)
-            if restored is None:
-                return None
-            restored.last_accessed_at = time.monotonic()
+            self._in_transit.add(dataset_id)
             try:
-                self.put(restored)
-            except CacheCapacityError:
-                # Spilled under a larger budget than this process has
-                # (e.g. TABCTX_GPU_MEMORY_FRACTION lowered across a
-                # restart). Nothing can re-admit it; treat as a miss so
-                # the caller re-fits instead of the request 500ing.
-                return None
-            return restored
+                restored = self._spill.load(dataset_id)
+                if restored is None:
+                    return None
+                restored.last_accessed_at = time.monotonic()
+                try:
+                    self.put(restored)
+                except CacheCapacityError:
+                    # Spilled under a larger budget than this process has
+                    # (e.g. TABCTX_GPU_MEMORY_FRACTION lowered across a
+                    # restart). Nothing can re-admit it; treat as a miss so
+                    # the caller re-fits instead of the request 500ing.
+                    return None
+                return restored
+            finally:
+                self._in_transit.discard(dataset_id)
 
     def touch(self, dataset_id: str) -> None:
         with self._lock:
@@ -110,9 +119,15 @@ def _evict_victim(self) -> str | None:
         victim_id = self._policy.select_victim(list(self._entries.values()))
         if victim_id is None:
             return None
-        victim = self._remove(victim_id)
-        if self._spill is not None:
+        if self._spill is None:
+            self._remove(victim_id)
+            return victim_id
+        self._in_transit.add(victim_id)
+        try:
+            victim = self._remove(victim_id)
             self._spill.spill(victim)  # best-effort by contract
+        finally:
+            self._in_transit.discard(victim_id)
         return victim_id
 
     def make_room(self, needed_bytes: int) -> list[str]:
@@ -181,9 +196,14 @@ def holds(self, dataset_id: str) -> bool:
         return dataset_id in self._entries
 
     def has(self, dataset_id: str) -> bool:
-        """Resident or spilled (snapshot, lock-free)."""
-        return dataset_id in self._entries or (
-            self._spill is not None and self._spill.contains(dataset_id)
+        """Resident, spilled, or moving between the two (snapshot,
+        lock-free). An id in transit is reported held: the caller's
+        locked get() waits for the move and finds it (or, if the spill
+        failed, misses then)."""
+        return (
+            dataset_id in self._entries
+            or dataset_id in self._in_transit
+            or (self._spill is not None and self._spill.contains(dataset_id))
         )
 
     def stats_nowait(self) -> EngineStats:
diff --git a/src/tabctx/multidevice.py b/src/tabctx/multidevice.py
index 33080cd..9f834c8 100644
--- a/src/tabctx/multidevice.py
+++ b/src/tabctx/multidevice.py
@@ -228,15 +228,23 @@ def _fit_on(self, index, X, y, task, dataset_id, backend, feature_names) -> str:
     def predict(
         self, dataset_id: str, X_test: ArrayLike, return_proba: bool = False
     ) -> PredictOutcome:
-        for index in self._locate(dataset_id):
-            try:
-                with self._active(index):
-                    return self._shards[index].predict(
-                        dataset_id, X_test, return_proba=return_proba
-                    )
-            except DatasetNotFoundError:
-                # Moved or evicted between the snapshot and the call.
-                continue
+        tried: set[int] = set()
+        # Two passes: a re-fit can move the id to another device between
+        # the first snapshot and the calls (it lands there, then the old
+        # copy is evicted), so a miss everywhere looks once more.
+        for _ in range(2):
+            for index in self._locate(dataset_id):
+                if index in tried:
+                    continue
+                tried.add(index)
+                try:
+                    with self._active(index):
+                        return self._shards[index].predict(
+                            dataset_id, X_test, return_proba=return_proba
+                        )
+                except DatasetNotFoundError:
+                    # Moved or evicted between the snapshot and the call.
+                    continue
         raise _not_found(dataset_id)
 
     def feature_names(self, dataset_id: str) -> list[str] | None:
diff --git a/tests/unit/test_multidevice.py b/tests/unit/test_multidevice.py
index 775e609..a3458eb 100644
--- a/tests/unit/test_multidevice.py
+++ b/tests/unit/test_multidevice.py
@@ -487,6 +487,54 @@ def test_concurrent_fits_spread_and_all_succeed(self):
             assert engine.predict(f"u{u}", [[0.0]]).predictions == ["a"]
 
 
+class TestMovingContexts:
+    """Regression (4xL4 rehearsal): a fit returned 200 and the next
+    predict 404'd within 30 ms -- the context was being spilled by a
+    concurrent fit's evict-ahead, so the router's lock-free lookup saw it
+    in neither tier."""
+
+    def test_predict_during_a_spill_waits_and_restores(self, tmp_path):
+        engine, shards, _ = build(2, spill_root=tmp_path)
+        engine.fit(*table(10, label="x"), dataset_id="d")
+        owner = where(shards, "d")[0]
+        cache = shards[owner]._cache
+        spill = cache._spill
+        gate, entered = threading.Event(), threading.Event()
+        original = spill.spill
+
+        def slow_spill(context):
+            entered.set()
+            assert gate.wait(5)
+            return original(context)
+
+        spill.spill = slow_spill
+        evicting, _ = start(cache.evict_one)
+        assert entered.wait(5)
+        predicting, box = start(engine.predict, "d", [[0.0]])
+        gate.set()
+        evicting.join(5)
+        predicting.join(5)
+        assert box.get("result") is not None, box.get("error")
+        assert box["result"].predictions == ["x"]
+
+    def test_predict_relocates_after_a_refit_moves_the_context(self):
+        engine, shards, _ = build(2)
+        engine.fit(*table(10, label="x"), dataset_id="d")
+        owner = where(shards, "d")[0]
+        stale = 1 - owner
+        real_locate = engine._locate
+        calls = []
+
+        def locate(dataset_id):
+            # First snapshot predates the move: only the old device.
+            calls.append(dataset_id)
+            return [stale] if len(calls) == 1 else real_locate(dataset_id)
+
+        engine._locate = locate
+        assert engine.predict("d", [[0.0]]).predictions == ["x"]
+        assert len(calls) == 2
+
+
 class TestFailures:
     def test_failed_refit_keeps_the_previous_context(self):
         engine, shards, _ = build(2)
diff --git a/tests/unit/test_spill.py b/tests/unit/test_spill.py
index 2905083..93ea73a 100644
--- a/tests/unit/test_spill.py
+++ b/tests/unit/test_spill.py
@@ -1,6 +1,7 @@
 """Unit tests for the disk spillover tier (cache/spill.py + its
 ContextCacheManager integration)."""
 
+import threading
 from pathlib import Path
 
 from tabctx.backends.fake import FakeBackend
@@ -201,3 +202,64 @@ def test_warm_restart_indexes_by_dataset_id(tmp_path):
     assert restarted.contains("t:alice:ds-1")
     restarted.delete("t:alice:ds-1")
     assert list(tmp_path.iterdir()) == []
+
+
+class GatedSpillStore(DiskSpillStore):
+    """Blocks inside spill()/load() until released, to observe the cache
+    while a context is between tiers."""
+
+    def __init__(self, directory) -> None:
+        super().__init__(directory)
+        self.gate = threading.Event()
+        self.entered = threading.Event()
+
+    def spill(self, context):
+        self.entered.set()
+        assert self.gate.wait(5), "spill gate never released"
+        return super().spill(context)
+
+    def load(self, dataset_id):
+        self.entered.set()
+        assert self.gate.wait(5), "load gate never released"
+        return super().load(dataset_id)
+
+
+class TestInTransitVisibility:
+    """A context being spilled or restored is listed by neither tier until
+    the I/O finishes; lock-free has() must still report it (a multi-device
+    router that didn't would answer 404 for a context that is moving)."""
+
+    def test_has_during_spill(self, tmp_path):
+        store = GatedSpillStore(tmp_path)
+        cache = ContextCacheManager(1000, spill_store=store)
+        cache.put(_ctx("moving"))
+        t = threading.Thread(target=cache.evict_one)
+        t.start()
+        assert store.entered.wait(5)
+        assert not cache.holds("moving")
+        assert cache.has("moving")
+        store.gate.set()
+        t.join(5)
+        assert cache.has("moving") and store.contains("moving")
+
+    def test_has_during_restore(self, tmp_path):
+        store = GatedSpillStore(tmp_path)
+        store.gate.set()
+        cache = ContextCacheManager(1000, spill_store=store)
+        cache.put(_ctx("moving"))
+        cache.evict_one()
+        store.gate.clear()
+        store.entered.clear()
+        t = threading.Thread(target=cache.get, args=("moving",))
+        t.start()
+        assert store.entered.wait(5)
+        assert cache.has("moving")
+        store.gate.set()
+        t.join(5)
+        assert cache.holds("moving")
+
+    def test_failed_spill_is_a_real_miss_afterwards(self, tmp_path):
+        cache = ContextCacheManager(1000, spill_store=DiskSpillStore(tmp_path))
+        cache.put(_ctx("bad", payload=lambda: None))
+        cache.evict_one()
+        assert not cache.has("bad")

From 2019511e13704adea80809cb18f84eed163edfe9 Mon Sep 17 00:00:00 2001
From: Amrit Krishnan <amrit110@gmail.com>
Date: Sat, 17 Oct 2026 18:06:59 -0400
Subject: [PATCH 10/39] bench: multi-GPU scaling benchmark (threads vs capped
 threads vs processes)

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 benchmarks/bench_scaling.py | 173 ++++++++++++++++++++++++++++++++++++
 benchmarks/gke/devbox.sh    |   2 +
 2 files changed, 175 insertions(+)
 create mode 100644 benchmarks/bench_scaling.py

diff --git a/benchmarks/bench_scaling.py b/benchmarks/bench_scaling.py
new file mode 100644
index 0000000..c56556a
--- /dev/null
+++ b/benchmarks/bench_scaling.py
@@ -0,0 +1,173 @@
+#!/usr/bin/env python3
+"""Does one process scale across a node's GPUs? (#39)
+
+The 4xL4 rehearsal ran the same 6-user workload on one replica using four
+GPUs and pinned to one: large fits took 142 s median on four GPUs vs 70 s
+on one. Concurrent fits on *different* GPUs slowed each other down. This
+isolates why, without Ray or HTTP: N workers each fit the same shape on
+their own GPU, all starting together, and report per-fit wall and CPU time.
+
+Modes:
+  threads         N threads in one process (how MultiDeviceEngine runs)
+  threads-capped  same, with torch's intra-op threads capped at cores // N
+  processes       one process per GPU, each capped at cores // N
+
+N=1 in any mode is the single-GPU baseline. If threads scale badly but
+processes scale well, the limit is in-process (GIL, a shared CPU thread
+pool); if both scale badly, it's the node (CPU cores, host memory or PCIe
+bandwidth).
+
+Usage: bench_scaling.py --backend tabicl --shape 40000x100 --gpus 1 2 4 \\
+           --modes threads threads-capped processes --out scaling.jsonl
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import multiprocessing as mp
+import os
+import resource
+import statistics
+import sys
+import threading
+import time
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+
+
+def _cores() -> int:
+    try:
+        return len(os.sched_getaffinity(0))
+    except AttributeError:
+        return os.cpu_count() or 1
+
+
+def _worker(kind, device, rows, features, fits, threads, start, results) -> None:
+    """Build a backend on `device`, wait for every worker, fit `fits` times."""
+    import torch
+    from profile_fit_predict import build_backend, make_table
+
+    if threads:
+        torch.set_num_threads(threads)
+    backend = build_backend(kind, device, "kv")
+    backend.warmup()
+    X, y = make_table(rows, features, "classification", seed=rows + features)
+    start.wait()
+    for _ in range(fits):
+        cpu0 = time.thread_time()
+        t0 = time.perf_counter()
+        payload = backend.fit(X, y, "classification")
+        torch.cuda.synchronize(device)
+        results.append(
+            {
+                "device": device,
+                "wall_s": time.perf_counter() - t0,
+                "thread_cpu_s": time.thread_time() - cpu0,
+            }
+        )
+        del payload
+
+
+def _process_entry(kind, device, rows, features, fits, threads, start, queue):
+    results: list = []
+    _worker(kind, device, rows, features, fits, threads, start, results)
+    for r in results:
+        queue.put(r)
+
+
+def run(kind: str, mode: str, n: int, rows: int, features: int, fits: int) -> dict:
+    devices = [f"cuda:{i}" for i in range(n)]
+    capped = max(1, _cores() // n) if mode != "threads" else 0
+    cpu0 = resource.getrusage(resource.RUSAGE_SELF)
+    t0 = time.perf_counter()
+    if mode == "processes":
+        ctx = mp.get_context("spawn")
+        start = ctx.Barrier(n)
+        queue = ctx.Queue()
+        procs = [
+            ctx.Process(
+                target=_process_entry,
+                args=(kind, d, rows, features, fits, capped, start, queue),
+            )
+            for d in devices
+        ]
+        for p in procs:
+            p.start()
+        results = [queue.get() for _ in range(n * fits)]
+        for p in procs:
+            p.join()
+        child = resource.getrusage(resource.RUSAGE_CHILDREN)
+        process_cpu = child.ru_utime + child.ru_stime
+    else:
+        start = threading.Barrier(n)
+        results = []
+        workers = [
+            threading.Thread(
+                target=_worker,
+                args=(kind, d, rows, features, fits, capped, start, results),
+            )
+            for d in devices
+        ]
+        for w in workers:
+            w.start()
+        for w in workers:
+            w.join()
+        now = resource.getrusage(resource.RUSAGE_SELF)
+        process_cpu = (now.ru_utime - cpu0.ru_utime) + (now.ru_stime - cpu0.ru_stime)
+    elapsed = time.perf_counter() - t0
+    walls = [r["wall_s"] for r in results]
+    return {
+        "backend": kind,
+        "mode": mode,
+        "gpus": n,
+        "rows": rows,
+        "features": features,
+        "fits": len(results),
+        "torch_threads": capped or None,
+        "cores": _cores(),
+        "fit_wall_p50_s": round(statistics.median(walls), 2),
+        "fit_wall_max_s": round(max(walls), 2),
+        "thread_cpu_p50_s": round(
+            statistics.median(r["thread_cpu_s"] for r in results), 2
+        ),
+        # Includes model loading/warmup; compare across modes, not to wall.
+        "process_cpu_s": round(process_cpu, 1),
+        "elapsed_s": round(elapsed, 1),
+    }
+
+
+def main() -> int:
+    p = argparse.ArgumentParser()
+    p.add_argument("--backend", choices=["tabicl", "tabpfn"], default="tabicl")
+    p.add_argument("--shape", default="40000x100")
+    p.add_argument("--gpus", type=int, nargs="+", default=[1, 2, 4])
+    p.add_argument(
+        "--modes",
+        nargs="+",
+        default=["threads", "threads-capped", "processes"],
+        choices=["threads", "threads-capped", "processes"],
+    )
+    p.add_argument("--fits", type=int, default=2, help="per worker")
+    p.add_argument("--out", default=None)
+    args = p.parse_args()
+    rows, features = (int(v) for v in args.shape.lower().split("x"))
+    out = None
+    if args.out:
+        Path(args.out).parent.mkdir(parents=True, exist_ok=True)
+        out = open(args.out, "a")  # noqa: SIM115
+    for mode in args.modes:
+        for n in args.gpus:
+            if mode != "threads" and n == 1 and "threads" in args.modes:
+                continue  # N=1 baseline is measured once
+            rec = run(args.backend, mode, n, rows, features, args.fits)
+            print(json.dumps(rec), flush=True)
+            if out:
+                out.write(json.dumps(rec) + "\n")
+                out.flush()
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/benchmarks/gke/devbox.sh b/benchmarks/gke/devbox.sh
index 0224105..075b63d 100755
--- a/benchmarks/gke/devbox.sh
+++ b/benchmarks/gke/devbox.sh
@@ -9,6 +9,8 @@
 #   benchmarks/gke/devbox.sh up            # create cluster + pod (MACHINES order)
 #   benchmarks/gke/devbox.sh sync          # copy src/ benchmarks/ tests/ into the pod
 #   benchmarks/gke/devbox.sh exec -- CMD   # run CMD in the pod (cwd /work/tabctx)
+#                                         # (CMD is joined into one shell string: quote
+#                                         #  pipes/redirects as one argument, no bash -c)
 #   benchmarks/gke/devbox.sh fetch PATH    # copy /work/tabctx/PATH back to ./gke-out/
 #   benchmarks/gke/devbox.sh down          # delete the cluster
 #

From 5a34568569c39687f1748dae4ab5e41a761ab109 Mon Sep 17 00:00:00 2001
From: Amrit Krishnan <amrit110@gmail.com>
Date: Sat, 17 Oct 2026 18:08:48 -0400
Subject: [PATCH 11/39] feat(memory): per-model requirement calibrations;
 width-aware resident slope

Requirement calibrations (floor points + requirement grid) are now looked
up per (model, cache mode) in a registry, so TabPFN gets its own measured
grid instead of TabICL's formula. ResidentRateModel extrapolates with the
per-row rate measured at the narrowest width >= the query's: TabPFN's
per-row cache grows with features, so a wider table's rate would
overshoot a narrow one's (TabICL's rate barely depends on width).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/tabctx/memory/requirement.py | 27 ++++++++++++++++++++++-----
 src/tabctx/serve/factory.py      | 30 ++++++++++++++++++++++--------
 tests/unit/test_requirement.py   | 20 ++++++++++++++++++++
 3 files changed, 64 insertions(+), 13 deletions(-)

diff --git a/src/tabctx/memory/requirement.py b/src/tabctx/memory/requirement.py
index 7fcfaa1..bf22b03 100644
--- a/src/tabctx/memory/requirement.py
+++ b/src/tabctx/memory/requirement.py
@@ -108,7 +108,10 @@ class ResidentRateModel:
     doesn't depend on features. So for any measured fit (n0 rows, d0 >= d
     features, resident r0), resident(n, d) <= r0 + rate * (n - n0) for
     n >= n0, where rate is the steepest per-row slope among long fits
-    (>= slope_min_rows; short fits are dominated by the column cache).
+    (>= slope_min_rows; short fits are dominated by the column cache) at
+    the narrowest measured width >= d: TabPFN's per-row cache grows with
+    features, so only a rate measured at least as wide bounds it, and the
+    narrowest such is the tightest.
     The estimate is the smallest such bound. Residents are fit-time sizes
     (before any post-fit compression -- what the fit needs while it runs),
     the same quantity the calibration measured. This replica's own long
@@ -132,14 +135,28 @@ def record(self, n_train: int, n_features: int, resident_bytes: int) -> None:
                 self._runtime.append((n_train, n_features, resident_bytes))
                 del self._runtime[:-500]  # bounded, newest kept
 
-    def _slope(self, points: list[tuple[int, int, int]]) -> float | None:
-        long_fits = [res / n for n, _, res in points if n >= self._min_rows]
-        return max(long_fits) if long_fits else None
+    def _slope(
+        self, points: list[tuple[int, int, int]], n_features: int
+    ) -> float | None:
+        """Steepest per-row rate among long fits at the narrowest measured
+        width >= n_features: residents don't shrink as features grow, so
+        that rate bounds this width's (and TabPFN's grows with width, so a
+        wider one would overshoot). Falls back to every long fit when none
+        is that wide (TabICL's rate barely depends on features)."""
+        long_fits = [(d, res / n) for n, d, res in points if n >= self._min_rows]
+        if not long_fits:
+            return None
+        widths = [d for d, _ in long_fits if d >= n_features]
+        if not widths:
+            return max(rate for _, rate in long_fits)
+        return max(rate for d, rate in long_fits if d == min(widths))
 
     def estimate(self, n_train: int, n_features: int) -> int | None:
         with self._lock:
             runtime = list(self._runtime)
-        slope = self._slope(runtime) or self._slope(self._calibration)
+        slope = self._slope(runtime, n_features) or self._slope(
+            self._calibration, n_features
+        )
         if slope is None:
             return None
         bounds = [
diff --git a/src/tabctx/serve/factory.py b/src/tabctx/serve/factory.py
index 7896d46..972887c 100644
--- a/src/tabctx/serve/factory.py
+++ b/src/tabctx/serve/factory.py
@@ -63,6 +63,7 @@
 
 from __future__ import annotations
 
+import importlib
 import logging
 import os
 import tempfile
@@ -297,21 +298,34 @@ def _preloaded_observations(
     return fit_grid, predict_grid
 
 
+# Generated requirement calibrations (benchmarks/gen_requirement_grid.py),
+# per (model, cache mode): module under tabctx.memory and its name prefix.
+_REQUIREMENT_CALIBRATIONS: dict[tuple[str, str], tuple[str, str]] = {
+    ("tabicl", "kv"): ("calibration_tabicl_a100_requirement", "A100_40GB_TABICL_KV"),
+    ("tabpfn", "kv"): ("calibration_tabpfn_a100_requirement", "A100_40GB_TABPFN_KV"),
+}
+
+
 def _requirement_calibration(
     kind: BackendKind, kv_cache: KvCacheMode
 ) -> tuple[WorkingFloorModel | None, tuple, tuple]:
-    """Measured working floors + requirements for models whose working
-    memory adapts to what is free (memory/requirement.py), so admission
-    bounds resident + floor instead of an unconstrained peak. Measured for
-    TabICL in kv mode so far; other configurations keep peak semantics."""
-    if kind != "tabicl" or kv_cache != "kv":
+    """Measured working floors + requirements (memory/requirement.py), so
+    admission bounds resident + floor instead of an unconstrained peak.
+    Configurations without a measured calibration keep peak semantics."""
+    entry = _REQUIREMENT_CALIBRATIONS.get((kind, kv_cache))
+    if entry is None:
         return None, (), ()
+    module_name, prefix = entry
     try:
-        from tabctx.memory import calibration_tabicl_a100_requirement as grid
+        grid = importlib.import_module(f"tabctx.memory.{module_name}")
     except ImportError:  # generated module absent: peak semantics, no preload
         return None, (), ()
-    points = grid.A100_40GB_TABICL_KV_FLOOR_POINTS
-    return WorkingFloorModel(points), grid.A100_40GB_TABICL_KV_REQUIREMENT_GRID, points
+    points = getattr(grid, f"{prefix}_FLOOR_POINTS")
+    return (
+        WorkingFloorModel(points),
+        getattr(grid, f"{prefix}_REQUIREMENT_GRID"),
+        points,
+    )
 
 
 def detect_gpu_capacity_bytes(devices: tuple[str, ...] = ("cuda",)) -> int | None:
diff --git a/tests/unit/test_requirement.py b/tests/unit/test_requirement.py
index 90885e9..2428cf3 100644
--- a/tests/unit/test_requirement.py
+++ b/tests/unit/test_requirement.py
@@ -160,6 +160,26 @@ def test_wider_query_uses_only_anchors_at_least_as_wide(self):
         )
         assert m.estimate(10, 5_000) is None
 
+    def test_narrow_tables_use_the_rate_measured_at_their_width(self):
+        from tabctx.memory.requirement import ResidentRateModel
+
+        # TabPFN-like: the per-row rate grows with features.
+        m = ResidentRateModel(
+            (
+                FloorPoint(50_000, 10, GiB, resident_bytes=1 * GiB),
+                FloorPoint(20_000, 100, GiB, resident_bytes=4 * GiB),
+            )
+        )
+        narrow = 1 * GiB / 50_000
+        assert m.estimate(100_000, 10) == pytest.approx(
+            1 * GiB + narrow * 50_000, rel=1e-6
+        )
+        # Wider than the narrow anchor: only the 100-feature rate bounds it.
+        wide = 4 * GiB / 20_000
+        assert m.estimate(40_000, 50) == pytest.approx(
+            4 * GiB + wide * 20_000, rel=1e-6
+        )
+
     def test_runtime_long_fits_supersede_the_calibration_slope(self):
         from tabctx.memory.requirement import ResidentRateModel
 

From f2c26716ae7a307196e6e4852877913e4ccc2a01 Mon Sep 17 00:00:00 2001
From: Amrit Krishnan <amrit110@gmail.com>
Date: Sat, 17 Oct 2026 18:09:15 -0400
Subject: [PATCH 12/39] bench: predict/fit overlap on one GPU (head-of-line
 study)

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 benchmarks/bench_overlap.py | 138 ++++++++++++++++++++++++++++++++++++
 1 file changed, 138 insertions(+)
 create mode 100644 benchmarks/bench_overlap.py

diff --git a/benchmarks/bench_overlap.py b/benchmarks/bench_overlap.py
new file mode 100644
index 0000000..93bc3ca
--- /dev/null
+++ b/benchmarks/bench_overlap.py
@@ -0,0 +1,138 @@
+#!/usr/bin/env python3
+"""Can predicts share a GPU with a running fit? (#39, head-of-line)
+
+Each device engine holds its cache lock across whole GPU calls, so a
+predict for a small context waits for any fit on that device to finish:
+on the 4xL4 rehearsal, predicts waited up to 266 s behind 40k x 100 fits.
+Interleaving them is only worth building if the GPU actually overlaps the
+work (a predict is a few small kernels; a fit is many large ones) and
+neither runs out of memory.
+
+On one device, with a small context already fitted, this measures:
+  1. predict latency on an idle device;
+  2. a large fit alone;
+  3. the same fit with a thread predicting the small context back to
+     back throughout (no lock between them): predict latency during the
+     fit, the fit's slowdown, and any OOM on either side.
+
+Usage: bench_overlap.py --device cuda:0 --fit 40000x100 --context 2000x20 \\
+           --out overlap.jsonl
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import statistics
+import sys
+import threading
+import time
+from pathlib import Path
+
+import torch
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+from profile_fit_predict import build_backend, make_table, parse_shape  # noqa: E402
+
+
+def _pcts(values: list[float]) -> dict:
+    if not values:
+        return {"n": 0}
+    ordered = sorted(values)
+    return {
+        "n": len(values),
+        "p50_ms": round(statistics.median(ordered) * 1000, 1),
+        "p95_ms": round(ordered[int(0.95 * (len(ordered) - 1))] * 1000, 1),
+        "max_ms": round(ordered[-1] * 1000, 1),
+    }
+
+
+def main() -> int:
+    p = argparse.ArgumentParser()
+    p.add_argument("--backend", choices=["tabicl", "tabpfn"], default="tabicl")
+    p.add_argument("--device", default="cuda:0")
+    p.add_argument("--fit", type=parse_shape, default=(40_000, 100))
+    p.add_argument("--context", type=parse_shape, default=(2_000, 20))
+    p.add_argument("--test-rows", type=int, default=100)
+    p.add_argument("--out", default=None)
+    args = p.parse_args()
+    device = torch.device(args.device)
+
+    backend = build_backend(args.backend, args.device, "kv")
+    backend.warmup()
+    cr, cf = args.context
+    Xc, yc = make_table(cr, cf, "classification", seed=1)
+    small = backend.fit(Xc, yc, "classification")
+    X_test, _ = make_table(args.test_rows, cf, "classification", seed=2)
+
+    def predict_once() -> float:
+        t0 = time.perf_counter()
+        backend.predict(small, X_test, False)
+        torch.cuda.synchronize(device)
+        return time.perf_counter() - t0
+
+    for _ in range(3):
+        predict_once()  # warm
+    idle = [predict_once() for _ in range(20)]
+
+    fr, ff = args.fit
+    Xf, yf = make_table(fr, ff, "classification", seed=3)
+
+    def timed_fit() -> tuple[float, str]:
+        torch.cuda.synchronize(device)
+        t0 = time.perf_counter()
+        try:
+            payload = backend.fit(Xf, yf, "classification")
+            torch.cuda.synchronize(device)
+            del payload
+            status = "ok"
+        except Exception as e:  # noqa: BLE001 -- an OOM is a result
+            status = f"{type(e).__name__}: {str(e)[:120]}"
+        torch.cuda.empty_cache()
+        return time.perf_counter() - t0, status
+
+    alone_s, alone_status = timed_fit()
+
+    during: list[float] = []
+    predict_errors: list[str] = []
+    stop = threading.Event()
+
+    def predictor() -> None:
+        while not stop.is_set():
+            try:
+                during.append(predict_once())
+            except Exception as e:  # noqa: BLE001
+                predict_errors.append(f"{type(e).__name__}: {str(e)[:120]}")
+                time.sleep(0.1)
+
+    t = threading.Thread(target=predictor, daemon=True)
+    t.start()
+    shared_s, shared_status = timed_fit()
+    stop.set()
+    t.join(30)
+
+    rec = {
+        "backend": args.backend,
+        "device": torch.cuda.get_device_name(device),
+        "fit": f"{fr}x{ff}",
+        "context": f"{cr}x{cf}",
+        "test_rows": args.test_rows,
+        "predict_idle": _pcts(idle),
+        "fit_alone_s": round(alone_s, 2),
+        "fit_alone_status": alone_status,
+        "fit_shared_s": round(shared_s, 2),
+        "fit_shared_status": shared_status,
+        "predict_during_fit": _pcts(during),
+        "predict_errors": predict_errors[:5],
+        "n_predict_errors": len(predict_errors),
+    }
+    print(json.dumps(rec), flush=True)
+    if args.out:
+        Path(args.out).parent.mkdir(parents=True, exist_ok=True)
+        with open(args.out, "a") as f:
+            f.write(json.dumps(rec) + "\n")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

From 6691b190cc847907ce38a3833fcb8b04c32ce8f6 Mon Sep 17 00:00:00 2001
From: Amrit Krishnan <amrit110@gmail.com>
Date: Sat, 17 Oct 2026 19:08:21 -0400
Subject: [PATCH 13/39] bench(pressure): fresh process per shape;
 --max-slowdown for models that degrade instead of failing

TabPFN falls back to slower paths under memory pressure rather than
raising, so 'the fit succeeded' put its floor near zero at 2-3x the fit
time. --max-slowdown counts such fits as failures. A long search process
also once measured a 5000x100 TabPFN fit at 55 s that takes 2-3 s in a
fresh process (not reproducible by replaying the constrained fits), so
each shape now runs in its own process.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 benchmarks/bench_memory_pressure.py | 47 +++++++++++++++++++++++------
 1 file changed, 38 insertions(+), 9 deletions(-)

diff --git a/benchmarks/bench_memory_pressure.py b/benchmarks/bench_memory_pressure.py
index 6294926..212608e 100644
--- a/benchmarks/bench_memory_pressure.py
+++ b/benchmarks/bench_memory_pressure.py
@@ -22,6 +22,7 @@
 
 import argparse
 import json
+import multiprocessing as mp
 import sys
 import time
 
@@ -77,10 +78,14 @@ def search(args, backend, rows: int, features: int) -> dict:
     lo = free_all["resident_gib"]  # can't succeed with less than the context
     hi = torch.cuda.mem_get_info(torch.device(args.device))[0] / GiB
     best = None
+    # A model that degrades instead of failing (TabPFN falls back to slower
+    # paths under memory pressure) "succeeds" at almost any free memory;
+    # its requirement is where it still runs at an acceptable speed.
+    max_s = free_all["fit_s"] * args.max_slowdown + 1.0 if args.max_slowdown else None
     while hi - lo > args.precision_gib:
         mid = (lo + hi) / 2
         r = run(args, backend, rows, features, round(mid, 2))
-        if r["status"] == "ok":
+        if r["status"] == "ok" and (max_s is None or r["fit_s"] <= max_s):
             hi, best = mid, r
         else:
             lo = mid
@@ -94,9 +99,21 @@ def search(args, backend, rows: int, features: int) -> dict:
         "min_free_gib": round(hi, 2),
         "min_working_gib": round(hi - free_all["resident_gib"], 2),
         "fit_s_at_min": best["fit_s"] if best else free_all["fit_s"],
+        "max_slowdown": args.max_slowdown,
     }
 
 
+def _measure_shape(args, rows: int, features: int, queue) -> None:
+    """One shape in a fresh process: no allocator or library state left
+    over from earlier shapes' OOMs can skew its measurements."""
+    backend = build_backend(args.backend, args.device, args.cache_mode)
+    backend.warmup()
+    if args.search:
+        queue.put([search(args, backend, rows, features)])
+    else:
+        queue.put([run(args, backend, rows, features, f) for f in args.free_gib])
+
+
 def main() -> int:
     p = argparse.ArgumentParser()
     p.add_argument("--backend", choices=["tabicl", "tabpfn"], required=True)
@@ -109,23 +126,35 @@ def main() -> int:
         "--search", action="store_true", help="binary-search min free memory"
     )
     p.add_argument("--precision-gib", type=float, default=0.5)
+    p.add_argument(
+        "--max-slowdown",
+        type=float,
+        default=None,
+        help="with --search: a fit slower than this x unconstrained counts as "
+        "failed (for models that degrade instead of failing, e.g. 2 for tabpfn)",
+    )
     p.add_argument("--out", default=None)
     args = p.parse_args()
 
-    backend = build_backend(args.backend, args.device, args.cache_mode)
-    backend.warmup()
-    out = open(args.out, "a") if args.out else None  # noqa: SIM115
+    ctx = mp.get_context("spawn")
     for rows, features in args.shapes:
+        queue = ctx.Queue()
+        child = ctx.Process(target=_measure_shape, args=(args, rows, features, queue))
+        child.start()
+        child.join()
         records = (
-            [search(args, backend, rows, features)]
-            if args.search
-            else [run(args, backend, rows, features, free) for free in args.free_gib]
+            queue.get()
+            if not queue.empty()
+            else [
+                {"rows": rows, "features": features, "status": f"exit {child.exitcode}"}
+            ]
         )
         for record in records:
             line = json.dumps(record)
             print(line, flush=True)
-            if out:
-                out.write(line + "\n")
+            if args.out:
+                with open(args.out, "a") as out:
+                    out.write(line + "\n")
     return 0
 
 

From 59b103a7988873ed31e8eddfa24ccf278b3b3665 Mon Sep 17 00:00:00 2001
From: Amrit Krishnan <amrit110@gmail.com>
Date: Sat, 17 Oct 2026 19:28:10 -0400
Subject: [PATCH 14/39] feat(serve): one worker process per GPU behind the
 multi-device router

Measured on 4x A100 (benchmarks/bench_scaling.py): TabICL fits are
host-bound (the driving thread is busy ~90% of wall time), so threads of
one process driving four GPUs contend for the GIL:

  per-fit wall      1 GPU    4 GPUs threads    4 GPUs processes
  10000 x 50        1.88 s   3.64 s (2.1x)     1.76 s (4.3x aggregate)
  40000 x 100      11.28 s   14.6 s (3.1x)     11.4 s (3.96x)

Capping torch's intra-op threads changed nothing; the GIL is the limit.

- serve/workers.py: ProcessShard proxies one device's TabctxEngine in a
  worker process (built by factory.build_shard). Calls are multiplexed
  over a pipe and served concurrently; each response carries a
  sequence-numbered snapshot (residency, stats, headroom) that answers
  the router's lock-free reads. A dead worker fails its calls with
  WorkerDiedError, is skipped by placement, and fails /readyz and Ray
  Serve's check_health so the replica restarts.
- multidevice.py: routes over a Shard protocol (TabctxEngine or
  ProcessShard); each shard prices a fit with its own estimator.
- factory.py: build_shard() shared by both modes; TABCTX_WORKER_PROCESSES
  (default 1) picks processes for multi-device replicas; one host-memory
  guard across processes (multiprocessing semaphore).
- app.py: real GPU memory and spill stats come from the workers (the
  serving process never creates CUDA contexts).

Nothing here depends on Ray: the router and workers are plain Python, so
the same core serves under Dynamo.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/tabctx/cache/manager.py      |  10 +
 src/tabctx/cache/spill.py        |   4 +
 src/tabctx/engine.py             |  11 +
 src/tabctx/memory/host.py        |  31 +-
 src/tabctx/multidevice.py        |  77 ++++-
 src/tabctx/serve/app.py          |  37 ++-
 src/tabctx/serve/factory.py      | 231 ++++++++++-----
 src/tabctx/serve/workers.py      | 476 +++++++++++++++++++++++++++++++
 tests/unit/test_serve_factory.py |   4 +-
 tests/unit/test_workers.py       | 128 +++++++++
 10 files changed, 916 insertions(+), 93 deletions(-)
 create mode 100644 src/tabctx/serve/workers.py
 create mode 100644 tests/unit/test_workers.py

diff --git a/src/tabctx/cache/manager.py b/src/tabctx/cache/manager.py
index 0b34302..c2bf4f5 100644
--- a/src/tabctx/cache/manager.py
+++ b/src/tabctx/cache/manager.py
@@ -206,6 +206,16 @@ def has(self, dataset_id: str) -> bool:
             or (self._spill is not None and self._spill.contains(dataset_id))
         )
 
+    def residency_nowait(self) -> tuple[frozenset[str], frozenset[str]]:
+        """(resident ids, every id has() reports: resident, in transit or
+        spilled) as lock-free snapshots -- what a router in another
+        process mirrors to answer holds()/has() without a round trip."""
+        resident = frozenset(list(self._entries))
+        known = resident | frozenset(list(self._in_transit))
+        if self._spill is not None:
+            known |= self._spill.ids()
+        return resident, known
+
     def stats_nowait(self) -> EngineStats:
         """stats() without the lock: fields may be from slightly different
         instants, which is fine for placement and monitoring."""
diff --git a/src/tabctx/cache/spill.py b/src/tabctx/cache/spill.py
index 48a6ba3..71ab405 100644
--- a/src/tabctx/cache/spill.py
+++ b/src/tabctx/cache/spill.py
@@ -160,6 +160,10 @@ def contains(self, dataset_id: str) -> bool:
         wait on another thread's spill I/O."""
         return dataset_id in self._index
 
+    def ids(self) -> frozenset[str]:
+        """Spilled ids (lock-free snapshot, like contains())."""
+        return frozenset(list(self._index))
+
     def delete(self, dataset_id: str) -> None:
         meta_path, payload_path = self._paths(dataset_id)
         payload_path.unlink(missing_ok=True)
diff --git a/src/tabctx/engine.py b/src/tabctx/engine.py
index 3f710b6..6fd73e1 100644
--- a/src/tabctx/engine.py
+++ b/src/tabctx/engine.py
@@ -191,6 +191,17 @@ def has(self, dataset_id: str) -> bool:
         """Resident or spilled here (lock-free snapshot)."""
         return self._cache.has(dataset_id)
 
+    def residency_nowait(self) -> tuple[frozenset[str], frozenset[str]]:
+        """(resident ids, resident-or-spilled ids), lock-free."""
+        return self._cache.residency_nowait()
+
+    def estimate_fit_bytes(
+        self, n_train: int, n_features: int, backend: str | None = None
+    ) -> int:
+        """What admission would charge a fit of this shape (placement uses
+        it to pick a device; the fit itself re-checks under the lock)."""
+        return self.estimator_for(backend).estimate_bytes(n_train, 0, n_features)
+
     def resolve_backend_name(self, backend: str | None) -> str:
         """The model a fit would use; InvalidInputError if unknown."""
         return self._resolve_backend(backend)[0]
diff --git a/src/tabctx/memory/host.py b/src/tabctx/memory/host.py
index fbaeb25..3aefaac 100644
--- a/src/tabctx/memory/host.py
+++ b/src/tabctx/memory/host.py
@@ -27,6 +27,7 @@
 import threading
 from collections.abc import Iterator
 from pathlib import Path
+from typing import Any
 
 GiB = 1024**3
 logger = logging.getLogger(__name__)
@@ -87,11 +88,17 @@ def _read_limit(limit_paths: tuple[str, ...], meminfo_path: str) -> int | None:
 
 
 class HostMemoryGuard:
+    """Process-wide by default. With mp_context, its semaphore and counters
+    are multiprocessing primitives, so one guard passed to per-GPU worker
+    processes (as a Process argument) bounds them all: host RAM is shared
+    by every process on the node."""
+
     def __init__(
         self,
         budget_bytes: int,
         large_fit_cells: int = LARGE_FIT_CELLS,
         large_fit_bytes: int = LARGE_FIT_HOST_BYTES,
+        mp_context: Any | None = None,
     ) -> None:
         self.slots = max(1, budget_bytes // large_fit_bytes)
         if budget_bytes < large_fit_bytes:
@@ -103,10 +110,14 @@ def __init__(
                 large_fit_bytes / GiB,
             )
         self._large_fit_cells = large_fit_cells
-        self._semaphore = threading.BoundedSemaphore(self.slots)
-        self._waiting = 0
-        self._active = 0
-        self._lock = threading.Lock()
+        if mp_context is None:
+            self._semaphore: Any = threading.BoundedSemaphore(self.slots)
+            self._lock: Any = threading.Lock()
+            self._counts: Any = [0, 0]  # waiting, active
+        else:
+            self._semaphore = mp_context.BoundedSemaphore(self.slots)
+            self._lock = mp_context.Lock()
+            self._counts = mp_context.Array("i", 2, lock=False)
 
     def is_large(self, n_rows: int, n_features: int) -> bool:
         """The fallback test for backends that can't say whether a call may
@@ -120,22 +131,22 @@ def reserve(self, large: bool) -> Iterator[None]:
             yield
             return
         with self._lock:
-            self._waiting += 1
+            self._counts[0] += 1
         self._semaphore.acquire()
         with self._lock:
-            self._waiting -= 1
-            self._active += 1
+            self._counts[0] -= 1
+            self._counts[1] += 1
         try:
             yield
         finally:
             with self._lock:
-                self._active -= 1
+                self._counts[1] -= 1
             self._semaphore.release()
 
     def stats(self) -> dict[str, int]:
         with self._lock:
             return {
                 "large_fit_slots": self.slots,
-                "large_fits_active": self._active,
-                "large_fits_waiting": self._waiting,
+                "large_fits_active": self._counts[1],
+                "large_fits_waiting": self._counts[0],
             }
diff --git a/src/tabctx/multidevice.py b/src/tabctx/multidevice.py
index 9f834c8..bd56b94 100644
--- a/src/tabctx/multidevice.py
+++ b/src/tabctx/multidevice.py
@@ -54,16 +54,62 @@
 import time
 from collections import deque
 from collections.abc import Callable, Iterator
+from typing import Protocol
 
-from tabctx.engine import TabctxEngine, _validate_fit_input
+from tabctx.engine import _validate_fit_input
 from tabctx.errors import DatasetNotFoundError
-from tabctx.memory.estimator import MemoryEstimator
 from tabctx.types import ArrayLike, EngineStats, PredictOutcome, Task
 
 # How long a started call counts toward a device's recent activity.
 ACTIVITY_WINDOW_S = 30.0
 
 
+class Shard(Protocol):
+    """One device's engine as the router sees it: a TabctxEngine in this
+    process, or a ProcessShard proxying one in a worker process
+    (serve/workers.py). The *_nowait-style reads (holds, has,
+    stats_nowait, admission_headroom_bytes) must not wait on the device's
+    GPU work."""
+
+    @property
+    def backend_names(self) -> list[str]: ...
+
+    @property
+    def default_backend(self) -> str: ...
+
+    def resolve_backend_name(self, backend: str | None) -> str: ...
+
+    def estimate_fit_bytes(
+        self, n_train: int, n_features: int, backend: str | None = None
+    ) -> int: ...
+
+    def admission_headroom_bytes(self, backend: str | None = None) -> int: ...
+
+    def holds(self, dataset_id: str) -> bool: ...
+
+    def has(self, dataset_id: str) -> bool: ...
+
+    def stats_nowait(self) -> EngineStats: ...
+
+    def fit(
+        self,
+        X: ArrayLike,
+        y: ArrayLike,
+        task: Task = "classification",
+        dataset_id: str | None = None,
+        backend: str | None = None,
+        feature_names: list[str] | None = None,
+    ) -> str: ...
+
+    def predict(
+        self, dataset_id: str, X_test: ArrayLike, return_proba: bool = False
+    ) -> PredictOutcome: ...
+
+    def feature_names(self, dataset_id: str) -> list[str] | None: ...
+
+    def evict(self, dataset_id: str) -> None: ...
+
+
 class _KeyedLocks:
     """One lock per key, created on demand and dropped when unused."""
 
@@ -96,7 +142,7 @@ def _not_found(dataset_id: str) -> DatasetNotFoundError:
 class MultiDeviceEngine:
     def __init__(
         self,
-        shards: list[TabctxEngine],
+        shards: list[Shard],
         devices: list[str],
         activity_window_s: float = ACTIVITY_WINDOW_S,
         clock: Callable[[], float] = time.monotonic,
@@ -128,9 +174,6 @@ def backend_names(self) -> list[str]:
     def default_backend(self) -> str:
         return self._shards[0].default_backend
 
-    def estimator_for(self, backend: str | None = None) -> MemoryEstimator:
-        return self._shards[0].estimator_for(backend)
-
     def admission_headroom_bytes(self, backend: str | None = None) -> int:
         """The best any one device offers: a fit lands on a single device."""
         return max(self.device_headroom_bytes(backend))
@@ -167,14 +210,19 @@ def _activity(self) -> list[tuple[int, int]]:
 
     # ---- placement and lookup (lock-free snapshots) ----------------
 
-    def _place(self, estimated_bytes: int, backend: str) -> int:
+    def _place(self, n_train: int, n_features: int, backend: str) -> int:
         headrooms = self.device_headroom_bytes(backend)
+        # Each shard prices the fit with its own estimator (worker
+        # processes learn separately; in-process shards share one).
+        needs = [
+            s.estimate_fit_bytes(n_train, n_features, backend) for s in self._shards
+        ]
         activity = self._activity()
         indices = range(len(self._shards))
-        fits = [i for i in indices if headrooms[i] >= estimated_bytes]
+        fits = [i for i in indices if headrooms[i] >= needs[i]]
         if fits:
             return min(fits, key=lambda i: (activity[i], -headrooms[i], i))
-        return max(indices, key=lambda i: (headrooms[i], -i))
+        return max(indices, key=lambda i: (headrooms[i] - needs[i], -i))
 
     def _locate(self, dataset_id: str) -> list[int]:
         """Devices holding the id: resident first, then spilled copies."""
@@ -199,13 +247,18 @@ def fit(
     ) -> str:
         n_features = _validate_fit_input(X, y)
         name = self._shards[0].resolve_backend_name(backend)
-        estimated = self.estimator_for(name).estimate_bytes(len(X), 0, n_features)
         if dataset_id is None:
             return self._fit_on(
-                self._place(estimated, name), X, y, task, None, name, feature_names
+                self._place(len(X), n_features, name),
+                X,
+                y,
+                task,
+                None,
+                name,
+                feature_names,
             )
         with self._id_locks.hold(dataset_id):
-            index = self._place(estimated, name)
+            index = self._place(len(X), n_features, name)
             resolved = self._fit_on(index, X, y, task, dataset_id, name, feature_names)
             # Only now that the new context is in place: drop copies on
             # other devices (a failed fit above leaves the old one serving).
diff --git a/src/tabctx/serve/app.py b/src/tabctx/serve/app.py
index a79a538..bf5a5da 100644
--- a/src/tabctx/serve/app.py
+++ b/src/tabctx/serve/app.py
@@ -213,9 +213,10 @@ def __init__(self) -> None:
         built = build_engine(settings)
         self._engine = built.engine
         self._estimator = built.estimator
-        self._backend_name = built.backend.name  # default model id
+        self._backend_name = built.default  # default model id
         self._spill_store = built.spill_store
         self._spill_stores = built.spill_stores
+        self._workers = built.worker_shards
         self._host_guard = built.host_guard
         self._devices = built.devices
         self._device = built.device
@@ -491,6 +492,14 @@ def _predict_sync(
     def healthz(self):
         return {"status": "ok"}
 
+    def check_health(self) -> None:
+        """Ray Serve's replica health probe: a dead per-GPU worker loses that
+        device's contexts and capacity for good, so fail the replica and
+        let Serve restart it whole."""
+        dead = [w.death_reason for w in self._workers if not w.alive]
+        if dead:
+            raise RuntimeError("; ".join(r for r in dead if r))
+
     @fastapi_app.get("/v1/models")
     def models(self):
         """Which models this deployment serves (chat-completions style
@@ -561,8 +570,21 @@ def readyz(self):
         torch = _torch_or_none()
 
         stats = self._engine.stats()
+        dead = [w.death_reason for w in self._workers if not w.alive]
+        if dead:
+            raise HTTPException(503, detail="; ".join(r for r in dead if r))
+        # Worker processes measure their own device (this process holds no
+        # CUDA context and must not create one on every GPU).
+        reports = [w.report() for w in self._workers]
         real_gpu_memory = None
-        if torch is not None and torch.cuda.is_available():
+        per_device: dict[str, dict] = {}
+        if reports:
+            per_device = {
+                w.device: r["gpu_memory"]
+                for w, r in zip(self._workers, reports, strict=True)
+                if r["gpu_memory"] is not None
+            }
+        elif torch is not None and torch.cuda.is_available():
             # Real device memory, independent of our own byte-accounting --
             # cache_stats.used_bytes is the estimator's *prediction* for what
             # cached contexts should cost; this is what the device actually
@@ -578,6 +600,7 @@ def readyz(self):
                 }
                 for d in self._cuda_devices()
             }
+        if per_device:
             # Same flat shape as ever (summed over devices when several),
             # plus the per-device breakdown on a multi-GPU replica.
             real_gpu_memory = {
@@ -606,7 +629,7 @@ def readyz(self):
                 "engine_calls": self._predictor.engine_calls,
             },
             "uploads": self._uploads.stats(),
-            "spill": self._spill_stats(),
+            "spill": self._spill_stats([r["spill"] for r in reports]),
             "real_gpu_memory": real_gpu_memory,
             "devices": self._device_breakdown(),
             "host_memory": self._host_guard.stats() if self._host_guard else None,
@@ -617,12 +640,14 @@ def _cuda_devices(self) -> list[str]:
         current device, on an auto single-GPU replica)."""
         return [d for d in self._devices if d.startswith("cuda")]
 
-    def _spill_stats(self) -> dict | None:
+    def _spill_stats(self, worker_spill: list[dict | None]) -> dict | None:
         if self._spill_store is not None:
             return self._spill_store.stats()
-        if not self._spill_stores:
+        per_device = [s.stats() for s in self._spill_stores] or [
+            s for s in worker_spill if s is not None
+        ]
+        if not per_device:
             return None
-        per_device = [s.stats() for s in self._spill_stores]
         totals: dict = {}
         for stats in per_device:
             for key, value in stats.items():
diff --git a/src/tabctx/serve/factory.py b/src/tabctx/serve/factory.py
index 972887c..271bbb4 100644
--- a/src/tabctx/serve/factory.py
+++ b/src/tabctx/serve/factory.py
@@ -59,6 +59,12 @@
   device. The replica only sees the GPUs its process is given: under Ray
   Serve that is ``ray_actor_options.num_gpus`` (or, for a ``num_gpus: 0``
   actor on a GPU pod, ``RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=1``).
+- ``TABCTX_WORKER_PROCESSES``: ``1`` (default) runs each device's engine
+  in its own worker process (serve/workers.py) when there are several
+  devices; ``0`` runs them as threads of the serving process. Fits are
+  host-bound, so threads contend for the GIL: measured on 4x A100, four
+  concurrent fits took 2x as long each from threads, single-GPU time
+  from processes.
 """
 
 from __future__ import annotations
@@ -102,6 +108,7 @@
 SPILL_CAPACITY_ENV_VAR = "TABCTX_SPILL_CAPACITY_BYTES"
 UPLOAD_TTL_S_ENV_VAR = "TABCTX_UPLOAD_TTL_S"
 DEVICES_ENV_VAR = "TABCTX_DEVICES"
+WORKER_PROCESSES_ENV_VAR = "TABCTX_WORKER_PROCESSES"
 
 logger = logging.getLogger(__name__)
 
@@ -126,6 +133,9 @@ class ServeSettings:
     devices: tuple[str, ...] = ()
     # None = <spill_dir>/offload, else a temporary directory.
     offload_dir: str | None = None
+    # Several devices: one worker process per GPU (default) or threads of
+    # this process (see _build_process_engine for why processes).
+    worker_processes: bool = True
 
     @classmethod
     def from_env(cls) -> ServeSettings:
@@ -202,6 +212,11 @@ def from_env(cls) -> ServeSettings:
             raise ValueError(f"{SPILL_CAPACITY_ENV_VAR} must be positive")
         devices = _parse_devices(os.environ.get(DEVICES_ENV_VAR, "auto"))
         offload_dir = os.environ.get(OFFLOAD_DIR_ENV_VAR) or None
+        raw_workers = os.environ.get(WORKER_PROCESSES_ENV_VAR, "1").strip().lower()
+        if raw_workers not in ("0", "1", "true", "false"):
+            raise ValueError(
+                f"{WORKER_PROCESSES_ENV_VAR}={raw_workers!r} must be 0 or 1"
+            )
         return cls(
             backends=backends,
             gpu_memory_fraction=fraction,
@@ -214,6 +229,7 @@ def from_env(cls) -> ServeSettings:
             spill_capacity_bytes=spill_capacity,
             devices=devices,
             offload_dir=offload_dir,
+            worker_processes=raw_workers in ("1", "true"),
         )
 
 
@@ -249,7 +265,7 @@ class BuiltEngine:
     engine: TabctxEngine | MultiDeviceEngine
     # Keyed by model/backend name; `default` names the first-listed model.
     estimators: dict[str, MemoryEstimator]
-    backends: dict[str, TabularICLBackend]
+    backend_names: list[str]
     default: str
     device: str
     spill_store: object | None = None
@@ -257,18 +273,22 @@ class BuiltEngine:
     # smallest when several); None = defaults.
     gpu_capacity_bytes: int | None = None
     devices: tuple[str, ...] = ()
-    # One per device when there are several (spill_store is then None).
+    # One per device when there are several in-process (spill_store is
+    # then None); worker processes report their own (worker_shards).
     spill_stores: tuple = ()
     host_guard: HostMemoryGuard | None = None
+    # Per-GPU worker processes (serve/workers.py), when used.
+    worker_shards: tuple = ()
 
-    # Single-model conveniences (the common deployment):
+    # Single-model convenience (the common deployment):
     @property
     def estimator(self) -> MemoryEstimator:
         return self.estimators[self.default]
 
-    @property
-    def backend(self) -> TabularICLBackend:
-        return self.backends[self.default]
+    def close(self) -> None:
+        """Stop worker processes, if any."""
+        for shard in self.worker_shards:
+            shard.close()
 
 
 def _preloaded_observations(
@@ -549,82 +569,165 @@ def _build_spill_store(
     return DiskSpillStore(path, capacity_bytes=capacity_bytes, serializers=serializers)
 
 
+@dataclass
+class Shard:
+    """One device's complete engine and the parts the serving layer reports
+    on (see build_shard)."""
+
+    engine: TabctxEngine
+    backends: dict[str, TabularICLBackend]
+    estimators: dict[str, MemoryEstimator]
+    spill_store: DiskSpillStore | None
+
+
+def build_shard(
+    settings: ServeSettings,
+    device: str,
+    *,
+    n_devices: int,
+    gpu_capacity: int | None,
+    host_guard: HostMemoryGuard | None,
+    estimators: dict[str, MemoryEstimator] | None = None,
+) -> Shard:
+    """Everything one device serves from: a backend per model bound to the
+    device, one cache (and spill tier) shared by those models, and the
+    engine over them. Used in-process and by per-GPU worker processes
+    (serve/workers.py) alike.
+
+    estimators: shared across shards of one process (identical cards
+    budgeted alike, so observations learned on one hold on the others);
+    filled in on first use. A worker process passes none and learns alone.
+    """
+    estimators = {} if estimators is None else estimators
+    fake_only = set(settings.backends) == {"fake"}
+    multi = n_devices > 1
+    backends: dict[str, TabularICLBackend] = {}
+    for kind in settings.backends:
+        backend = _build_backend(
+            kind, settings, device, host_guard.slots if host_guard else 1
+        )
+        backends[backend.name] = backend
+        if backend.name not in estimators:
+            estimators[backend.name] = build_estimator(settings, kind, gpu_capacity)
+    # Backends are keyed by model id (e.g. "tabicl-v2"), not by kind
+    # ("tabicl"); insertion order follows settings.backends, so the
+    # first-listed kind's model is the default.
+    default = next(iter(backends))
+    spill_store = _build_spill_store(
+        settings,
+        backends,
+        device.replace(":", "_") if multi else None,
+        settings.spill_capacity_bytes // n_devices,
+    )
+    # ONE cache and budget shared by every model on a device -- the whole
+    # point of serving them behind one endpoint (see engine.py).
+    cache = ContextCacheManager(
+        capacity_bytes=estimators[default].ceiling_bytes(), spill_store=spill_store
+    )
+    engine = TabctxEngine(
+        backends=backends,
+        estimators={name: estimators[name] for name in backends},
+        cache=cache,
+        default_backend=default,
+        host_guard=host_guard,
+        device_baseline_bytes=0 if fake_only else measure_device_baseline(device),
+    )
+    return Shard(
+        engine=engine,
+        backends=backends,
+        estimators={name: estimators[name] for name in backends},
+        spill_store=spill_store,
+    )
+
+
 def build_engine(settings: ServeSettings | None = None) -> BuiltEngine:
     settings = settings or ServeSettings.from_env()
     devices = resolve_devices(settings)
+    fake_only = set(settings.backends) == {"fake"}
     # Budget against the real device, not the calibration card's size;
     # the fake backend never touches a GPU even when one is visible.
-    gpu_capacity = (
-        None
-        if set(settings.backends) == {"fake"}
-        else detect_gpu_capacity_bytes(devices)
-    )
+    gpu_capacity = None if fake_only else detect_gpu_capacity_bytes(devices)
     multi = len(devices) > 1
-    fake_only = set(settings.backends) == {"fake"}
     if settings.spill_dir:
         _warn_on_spill_layout(settings.spill_dir, multi)
-    estimators: dict[str, MemoryEstimator] = {}
+    host_budget = None if fake_only else detect_host_budget_bytes()
+    if multi and settings.worker_processes:
+        return _build_process_engine(settings, devices, gpu_capacity, host_budget)
     # One host-memory guard per process: host RAM is shared by every device.
-    host_budget = (
-        None if set(settings.backends) == {"fake"} else detect_host_budget_bytes()
-    )
     host_guard = HostMemoryGuard(host_budget) if host_budget else None
-    shards: list[TabctxEngine] = []
-    spill_stores: list[DiskSpillStore | None] = []
-    first_backends: dict[str, TabularICLBackend] = {}
-    for device in devices:
-        backends: dict[str, TabularICLBackend] = {}
-        for kind in settings.backends:
-            backend = _build_backend(
-                kind, settings, device, host_guard.slots if host_guard else 1
-            )
-            backends[backend.name] = backend
-            # One estimator per model, shared by every device: they are
-            # identical cards budgeted alike, and observations learned on
-            # one device are valid on the others.
-            if backend.name not in estimators:
-                estimators[backend.name] = build_estimator(settings, kind, gpu_capacity)
-        # Backends are keyed by model id (e.g. "tabicl-v2"), not by kind
-        # ("tabicl"); insertion order follows settings.backends, so the
-        # first-listed kind's model is the default.
-        default = next(iter(backends))
-        spill_store = _build_spill_store(
+    estimators: dict[str, MemoryEstimator] = {}
+    shards = [
+        build_shard(
             settings,
-            backends,
-            device.replace(":", "_") if multi else None,
-            settings.spill_capacity_bytes // len(devices),
-        )
-        # ONE cache and budget shared by every model on a device -- the
-        # whole point of serving them behind one endpoint (see engine.py).
-        cache = ContextCacheManager(
-            capacity_bytes=estimators[default].ceiling_bytes(), spill_store=spill_store
+            device,
+            n_devices=len(devices),
+            gpu_capacity=gpu_capacity,
+            host_guard=host_guard,
+            estimators=estimators,
         )
-        shards.append(
-            TabctxEngine(
-                backends=backends,
-                estimators={name: estimators[name] for name in backends},
-                cache=cache,
-                default_backend=default,
-                host_guard=host_guard,
-                device_baseline_bytes=0
-                if fake_only
-                else measure_device_baseline(device),
-            )
-        )
-        spill_stores.append(spill_store)
-        first_backends = first_backends or backends
+        for device in devices
+    ]
+    first = shards[0]
     engine: TabctxEngine | MultiDeviceEngine = (
-        MultiDeviceEngine(shards, list(devices)) if multi else shards[0]
+        MultiDeviceEngine([s.engine for s in shards], list(devices))
+        if multi
+        else first.engine
     )
     return BuiltEngine(
         engine=engine,
-        estimators={name: estimators[name] for name in first_backends},
-        backends=first_backends,
-        default=next(iter(first_backends)),
+        estimators=first.estimators,
+        backend_names=list(first.backends),
+        default=next(iter(first.backends)),
         device="cpu (fake backend)" if fake_only else ",".join(devices),
-        spill_store=None if multi else spill_stores[0],
+        spill_store=None if multi else first.spill_store,
+        gpu_capacity_bytes=gpu_capacity,
+        devices=devices,
+        spill_stores=tuple(s.spill_store for s in shards if s.spill_store is not None)
+        if multi
+        else (),
+        host_guard=host_guard,
+    )
+
+
+def _build_process_engine(
+    settings: ServeSettings,
+    devices: tuple[str, ...],
+    gpu_capacity: int | None,
+    host_budget: int | None,
+) -> BuiltEngine:
+    """One worker process per GPU (serve/workers.py), routed by a
+    MultiDeviceEngine in this process. Measured on 4x A100 (benchmarks/
+    bench_scaling.py): concurrent fits on four GPUs from threads of one
+    process took 2x as long each (the GIL -- fits are host-bound), while
+    a process per GPU ran each at single-GPU speed."""
+    import multiprocessing as mp
+
+    from tabctx.serve.workers import ProcessShard
+
+    ctx = mp.get_context("spawn")
+    host_guard = HostMemoryGuard(host_budget, mp_context=ctx) if host_budget else None
+    shards = ProcessShard.start_all(
+        settings,
+        devices,
+        gpu_capacity=gpu_capacity,
+        host_guard=host_guard,
+        mp_context=ctx,
+    )
+    # This process holds no backends: /limits and /readyz use estimators
+    # built from the same settings (calibration, not this replica's
+    # learned observations).
+    estimators = {
+        name: build_estimator(settings, kind, gpu_capacity)
+        for name, kind in zip(shards[0].backend_names, settings.backends, strict=True)
+    }
+    return BuiltEngine(
+        engine=MultiDeviceEngine(list(shards), list(devices)),
+        estimators=estimators,
+        backend_names=list(shards[0].backend_names),
+        default=shards[0].default_backend,
+        device=",".join(devices),
         gpu_capacity_bytes=gpu_capacity,
         devices=devices,
-        spill_stores=tuple(s for s in spill_stores if s is not None) if multi else (),
         host_guard=host_guard,
+        worker_shards=tuple(shards),
     )
diff --git a/src/tabctx/serve/workers.py b/src/tabctx/serve/workers.py
new file mode 100644
index 0000000..412fe64
--- /dev/null
+++ b/src/tabctx/serve/workers.py
@@ -0,0 +1,476 @@
+"""Per-GPU worker processes: one TabctxEngine per device, each in its own
+process, routed by a MultiDeviceEngine in the serving process.
+
+Why processes. TabICL and TabPFN fits are host-bound: the Python thread
+driving a fit is busy for ~90% of its wall time (kernel launches,
+preprocessing). Threads of one process driving several GPUs therefore
+contend for the GIL -- measured on 4x A100 (benchmarks/bench_scaling.py),
+four concurrent 10k x 50 fits took 3.6 s each from threads vs 1.9 s
+alone, and 1.8 s each from a process per GPU (4.3x aggregate vs 2.1x).
+
+Shape. ``ProcessShard`` implements the router's ``Shard`` protocol
+(multidevice.py) by proxying calls to a worker process that owns one
+device's complete engine (built by factory.build_shard: backends, cache,
+spill tier, its own estimator). The serving layer, the router and its
+placement policy are unchanged, and none of this depends on Ray: the
+same core runs under any serving framework that hosts one Python process
+per node (Ray Serve today, Dynamo tomorrow).
+
+Protocol. One duplex pipe per worker; requests ``(id, method, args,
+kwargs)`` are multiplexed (a worker serves several at once from a thread
+pool, so predicts queue behind a fit exactly as in-process), and every
+response carries a ``Snapshot`` of the worker's residency, stats and
+admission headroom. The router's lock-free reads (holds, has,
+stats_nowait, admission_headroom_bytes) answer from the latest snapshot,
+never a round trip. Snapshots are sequence-numbered so a late response
+can't roll the view back. The view may lag the worker's truth, safely:
+a context the view still lists but the worker evicted makes a predict
+miss there and the router look elsewhere (a real 404 only if nowhere);
+a new context is listed before its fit's response returns to the caller.
+
+Failure. If a worker dies (host OOM kill, a CUDA fatal error), its
+pending calls fail with WorkerDiedError (a BackendComputeError, so the
+API answers 500, not a hang), the shard reports itself dead, placement
+skips it, and /readyz reports the replica unhealthy so the serving
+framework restarts it.
+"""
+
+from __future__ import annotations
+
+import itertools
+import logging
+import pickle
+import threading
+import time
+from concurrent.futures import Future, ThreadPoolExecutor
+from dataclasses import dataclass
+from typing import Any
+
+from tabctx.errors import BackendComputeError
+from tabctx.types import ArrayLike, EngineStats, PredictOutcome, Task
+
+logger = logging.getLogger(__name__)
+
+# Calls the worker serves (the Shard surface's blocking part, plus reports).
+_METHODS = frozenset(
+    {
+        "fit",
+        "predict",
+        "evict",
+        "feature_names",
+        "estimate_fit_bytes",
+        "report",
+    }
+)
+# Concurrent calls per worker: its engine serializes GPU work itself
+# (cache lock); these only need to queue and answer cheap calls promptly.
+_WORKER_THREADS = 16
+
+
+class WorkerDiedError(BackendComputeError):
+    """A worker process exited with calls in flight."""
+
+
+@dataclass(frozen=True)
+class Snapshot:
+    seq: int
+    resident: frozenset[str]
+    known: frozenset[str]  # resident, in transit or spilled
+    stats: EngineStats
+    headroom: dict[str, int]  # backend -> admission headroom bytes
+
+
+@dataclass(frozen=True)
+class _Ready:
+    backend_names: list[str]
+    default_backend: str
+    snapshot: Snapshot
+
+
+@dataclass
+class _Failure:
+    error: BaseException
+
+
+def _portable(error: BaseException) -> BaseException:
+    """The exception itself if it pickles (tabctx errors do), else a
+    RuntimeError carrying its type and message."""
+    try:
+        pickle.dumps(error)
+        return error
+    except Exception:  # noqa: BLE001
+        return RuntimeError(f"{type(error).__name__}: {error}")
+
+
+# ---- worker process --------------------------------------------------------
+
+
+def _snapshot(engine: Any, seq: int) -> Snapshot:
+    resident, known = engine.residency_nowait()
+    return Snapshot(
+        seq=seq,
+        resident=resident,
+        known=known,
+        stats=engine.stats_nowait(),
+        headroom={b: engine.admission_headroom_bytes(b) for b in engine.backend_names},
+    )
+
+
+def _torch_if_cuda() -> Any:
+    try:
+        import torch
+    except ImportError:
+        return None
+    return torch if torch.cuda.is_available() else None
+
+
+def _report(shard: Any, device: str) -> dict:
+    """What /readyz shows per device that only the owning process can
+    measure: its spill tier and the GPU memory torch really holds."""
+    gpu = None
+    torch = _torch_if_cuda() if device.startswith("cuda") else None
+    if torch is not None:
+        gpu = {
+            "allocated_mb": torch.cuda.memory_allocated(device) / 1e6,
+            "reserved_mb": torch.cuda.memory_reserved(device) / 1e6,
+        }
+    return {
+        "spill": shard.spill_store.stats() if shard.spill_store else None,
+        "gpu_memory": gpu,
+    }
+
+
+def _worker_main(
+    conn: Any,
+    settings: Any,
+    device: str,
+    n_devices: int,
+    gpu_capacity: int | None,
+    host_guard: Any,
+) -> None:
+    """Entry point of a worker process: build the device's engine, then
+    serve calls until the pipe closes."""
+    logging.basicConfig(level=logging.INFO)
+    from tabctx.serve.factory import build_shard
+
+    try:
+        shard = build_shard(
+            settings,
+            device,
+            n_devices=n_devices,
+            gpu_capacity=gpu_capacity,
+            host_guard=host_guard,
+        )
+    except BaseException as e:  # noqa: BLE001 -- reported to the parent
+        conn.send(_Failure(_portable(e)))
+        return
+    engine = shard.engine
+    seq = itertools.count(1)
+    seq_lock = threading.Lock()
+    send_lock = threading.Lock()
+
+    def snapshot() -> Snapshot:
+        # Numbered and taken under one lock, so a higher seq is never an
+        # older view.
+        with seq_lock:
+            return _snapshot(engine, next(seq))
+
+    def handle(req_id: int, method: str, args: tuple, kwargs: dict) -> None:
+        try:
+            if method == "report":
+                result: Any = _report(shard, device)
+            else:
+                result = getattr(engine, method)(*args, **kwargs)
+            reply = (req_id, True, result, snapshot())
+        except BaseException as e:  # noqa: BLE001 -- returned to the caller
+            reply = (req_id, False, _portable(e), snapshot())
+        with send_lock:
+            conn.send(reply)
+
+    conn.send(_Ready(list(engine.backend_names), engine.default_backend, snapshot()))
+    with ThreadPoolExecutor(_WORKER_THREADS, thread_name_prefix=f"w-{device}") as pool:
+        while True:
+            try:
+                msg = conn.recv()
+            except (EOFError, OSError):
+                break
+            if msg is None:  # orderly shutdown
+                break
+            req_id, method, args, kwargs = msg
+            if method not in _METHODS:
+                with send_lock:
+                    conn.send(
+                        (
+                            req_id,
+                            False,
+                            ValueError(f"unknown call {method!r}"),
+                            snapshot(),
+                        )
+                    )
+                continue
+            pool.submit(handle, req_id, method, args, kwargs)
+
+
+# ---- serving process -------------------------------------------------------
+
+
+class ProcessShard:
+    """The router's Shard for one device whose engine runs in a worker
+    process. Thread-safe: any number of serving threads may call it."""
+
+    def __init__(
+        self,
+        device: str,
+        settings: Any,
+        *,
+        n_devices: int,
+        gpu_capacity: int | None,
+        host_guard: Any,
+        mp_context: Any,
+    ) -> None:
+        self.device = device
+        parent, child = mp_context.Pipe()
+        self._conn = parent
+        self._proc = mp_context.Process(
+            target=_worker_main,
+            args=(child, settings, device, n_devices, gpu_capacity, host_guard),
+            name=f"tabctx-worker-{device}",
+            daemon=True,
+        )
+        self._proc.start()
+        child.close()
+        self._send_lock = threading.Lock()
+        self._state_lock = threading.Lock()
+        self._pending: dict[int, Future] = {}
+        self._ids = itertools.count(1)
+        self._dead: str | None = None
+        self._snapshot: Snapshot | None = None
+        self._backend_names: list[str] = []
+        self._default_backend = ""
+        self._reader: threading.Thread | None = None
+
+    # ---- lifecycle ---------------------------------------------------
+
+    def wait_ready(self, timeout_s: float) -> None:
+        """Block until the worker built its engine (loading model weights);
+        raise its error if it failed."""
+        if not self._conn.poll(timeout_s):
+            self.close()
+            raise TimeoutError(f"worker for {self.device} not ready in {timeout_s}s")
+        try:
+            first = self._conn.recv()
+        except (EOFError, OSError) as e:
+            raise WorkerDiedError(
+                f"worker for {self.device} exited during startup "
+                f"(exit code {self._proc.exitcode})"
+            ) from e
+        if isinstance(first, _Failure):
+            self._proc.join(5)
+            raise first.error
+        self._backend_names = first.backend_names
+        self._default_backend = first.default_backend
+        self._snapshot = first.snapshot
+        self._reader = threading.Thread(
+            target=self._read_loop, name=f"tabctx-reader-{self.device}", daemon=True
+        )
+        self._reader.start()
+
+    @classmethod
+    def start_all(
+        cls,
+        settings: Any,
+        devices: tuple[str, ...],
+        *,
+        gpu_capacity: int | None,
+        host_guard: Any,
+        mp_context: Any,
+        ready_timeout_s: float = 900.0,
+    ) -> list[ProcessShard]:
+        """Start one worker per device in parallel (each loads its own model
+        weights) and wait for all of them."""
+        shards: list[ProcessShard] = [
+            cls(
+                d,
+                settings,
+                n_devices=len(devices),
+                gpu_capacity=gpu_capacity,
+                host_guard=host_guard,
+                mp_context=mp_context,
+            )
+            for d in devices
+        ]
+        try:
+            deadline = time.monotonic() + ready_timeout_s
+            for s in shards:
+                s.wait_ready(max(1.0, deadline - time.monotonic()))
+        except BaseException:
+            for s in shards:
+                s.close()
+            raise
+        return shards
+
+    def close(self) -> None:
+        with self._state_lock:
+            if self._dead is None:
+                self._dead = f"worker for {self.device} was shut down"
+        with self._send_lock:
+            try:
+                self._conn.send(None)
+            except (OSError, ValueError):
+                pass
+        self._proc.join(10)
+        if self._proc.is_alive():
+            self._proc.kill()
+            self._proc.join(5)
+        self._conn.close()
+        self._fail_pending(f"worker for {self.device} was shut down")
+
+    @property
+    def alive(self) -> bool:
+        return self._dead is None and self._proc.is_alive()
+
+    @property
+    def death_reason(self) -> str | None:
+        if self._dead is None and not self._proc.is_alive():
+            return f"worker for {self.device} exited (code {self._proc.exitcode})"
+        return self._dead
+
+    # ---- transport ---------------------------------------------------
+
+    def _read_loop(self) -> None:
+        while True:
+            try:
+                req_id, ok, payload, snap = self._conn.recv()
+            except (EOFError, OSError):
+                break
+            self._apply(snap)
+            with self._state_lock:
+                fut = self._pending.pop(req_id, None)
+            if fut is None:
+                continue
+            if ok:
+                fut.set_result(payload)
+            else:
+                fut.set_exception(payload)
+        self._proc.join(1)  # reap it, for a meaningful exit code
+        reason = (
+            f"worker for {self.device} exited (code {self._proc.exitcode}); "
+            "its contexts are lost"
+        )
+        if self._dead is None:
+            logger.error(reason)
+        self._fail_pending(reason)
+
+    def _fail_pending(self, reason: str) -> None:
+        with self._state_lock:
+            if self._dead is None:
+                self._dead = reason
+            pending, self._pending = self._pending, {}
+        for fut in pending.values():
+            if not fut.done():
+                fut.set_exception(WorkerDiedError(reason))
+
+    def _apply(self, snap: Snapshot) -> None:
+        with self._state_lock:
+            if self._snapshot is None or snap.seq > self._snapshot.seq:
+                self._snapshot = snap
+
+    def _call(self, method: str, *args: Any, **kwargs: Any) -> Any:
+        fut: Future = Future()
+        with self._state_lock:
+            if self._dead is not None:
+                raise WorkerDiedError(self._dead)
+            req_id = next(self._ids)
+            self._pending[req_id] = fut
+        try:
+            with self._send_lock:
+                self._conn.send((req_id, method, args, kwargs))
+        except (OSError, ValueError) as e:
+            with self._state_lock:
+                self._pending.pop(req_id, None)
+            raise WorkerDiedError(f"worker for {self.device} is gone: {e}") from e
+        return fut.result()
+
+    # ---- Shard: lock-free reads from the latest snapshot -------------
+
+    def _snap(self) -> Snapshot:
+        snap = self._snapshot
+        assert snap is not None, "wait_ready() first"
+        return snap
+
+    @property
+    def backend_names(self) -> list[str]:
+        return list(self._backend_names)
+
+    @property
+    def default_backend(self) -> str:
+        return self._default_backend
+
+    def resolve_backend_name(self, backend: str | None) -> str:
+        from tabctx.errors import InvalidInputError
+
+        name = backend or self._default_backend
+        if name not in self._backend_names:
+            raise InvalidInputError(
+                f"unknown backend {name!r}; this deployment serves "
+                f"{', '.join(self._backend_names)}"
+            )
+        return name
+
+    def holds(self, dataset_id: str) -> bool:
+        return self.alive and dataset_id in self._snap().resident
+
+    def has(self, dataset_id: str) -> bool:
+        return self.alive and dataset_id in self._snap().known
+
+    def stats_nowait(self) -> EngineStats:
+        return self._snap().stats
+
+    def admission_headroom_bytes(self, backend: str | None = None) -> int:
+        # A dead worker offers nothing, so placement never picks it.
+        if not self.alive:
+            return 0
+        return self._snap().headroom[backend or self._default_backend]
+
+    # ---- Shard: calls ------------------------------------------------
+
+    def estimate_fit_bytes(
+        self, n_train: int, n_features: int, backend: str | None = None
+    ) -> int:
+        if not self.alive:
+            return 1 << 62  # never fits: placement skips a dead worker
+        return self._call("estimate_fit_bytes", n_train, n_features, backend)
+
+    def fit(
+        self,
+        X: ArrayLike,
+        y: ArrayLike,
+        task: Task = "classification",
+        dataset_id: str | None = None,
+        backend: str | None = None,
+        feature_names: list[str] | None = None,
+    ) -> str:
+        return self._call(
+            "fit",
+            X,
+            y,
+            task=task,
+            dataset_id=dataset_id,
+            backend=backend,
+            feature_names=feature_names,
+        )
+
+    def predict(
+        self, dataset_id: str, X_test: ArrayLike, return_proba: bool = False
+    ) -> PredictOutcome:
+        return self._call("predict", dataset_id, X_test, return_proba=return_proba)
+
+    def feature_names(self, dataset_id: str) -> list[str] | None:
+        return self._call("feature_names", dataset_id)
+
+    def evict(self, dataset_id: str) -> None:
+        self._call("evict", dataset_id)
+
+    def report(self) -> dict:
+        """The worker's spill stats and real GPU memory (see _report)."""
+        return self._call("report")
diff --git a/tests/unit/test_serve_factory.py b/tests/unit/test_serve_factory.py
index c6473eb..0d99987 100644
--- a/tests/unit/test_serve_factory.py
+++ b/tests/unit/test_serve_factory.py
@@ -208,7 +208,7 @@ def test_detect_returns_none_without_a_gpu(self, monkeypatch):
 class TestBuildEngine:
     def test_fake_backend_end_to_end(self):
         built = build_engine(ServeSettings(backends=("fake",)))
-        assert isinstance(built.backend, FakeBackend)
+        assert built.backend_names == ["fake"] and built.default == "fake"
         assert "fake" in built.device
 
         dataset_id = built.engine.fit([[1.0, 2.0], [3.0, 4.0]], ["a", "b"])
@@ -402,6 +402,7 @@ def test_several_devices_build_one_engine_each(self, monkeypatch, tmp_path):
                 devices=("cuda:0", "cuda:1", "cuda:2"),
                 spill_dir=str(tmp_path),
                 spill_capacity_bytes=30 * 1024**3,
+                worker_processes=False,  # the patches above are in-process
             )
         )
         assert isinstance(built.engine, MultiDeviceEngine)
@@ -455,6 +456,7 @@ def test_spill_layout_change_is_reported(self, tmp_path, caplog):
                     backends=("fake",),
                     devices=("cuda:0", "cuda:1"),
                     spill_dir=str(tmp_path),
+                    worker_processes=False,
                 )
             )
         assert "will not be restored" in caplog.text
diff --git a/tests/unit/test_workers.py b/tests/unit/test_workers.py
new file mode 100644
index 0000000..c1cacb6
--- /dev/null
+++ b/tests/unit/test_workers.py
@@ -0,0 +1,128 @@
+"""Per-GPU worker processes (serve/workers.py), with real processes and the
+fake backend (which never touches a GPU, so 'cuda:N' are just labels)."""
+
+import threading
+
+import pytest
+
+from tabctx.errors import DatasetNotFoundError, InvalidInputError
+from tabctx.multidevice import MultiDeviceEngine
+from tabctx.serve.factory import ServeSettings, build_engine
+from tabctx.serve.workers import ProcessShard, WorkerDiedError
+
+
+def table(rows: int, label: str = "a", features: int = 2):
+    return [[float(i)] * features for i in range(rows)], [label] * rows
+
+
+def _settings(tmp_path=None, **kw) -> ServeSettings:
+    return ServeSettings(
+        backends=("fake",),
+        devices=("cuda:0", "cuda:1"),
+        spill_dir=str(tmp_path) if tmp_path else None,
+        **kw,
+    )
+
+
+@pytest.fixture(scope="module")
+def built(tmp_path_factory):
+    b = build_engine(_settings(tmp_path_factory.mktemp("spill")))
+    yield b
+    b.close()
+
+
+def test_multi_device_defaults_to_worker_processes(built):
+    assert isinstance(built.engine, MultiDeviceEngine)
+    assert [type(w) for w in built.worker_shards] == [ProcessShard, ProcessShard]
+    assert all(w.alive for w in built.worker_shards)
+    assert built.backend_names == ["fake"] and built.default == "fake"
+
+
+def test_threads_when_disabled():
+    b = build_engine(_settings(worker_processes=False))
+    assert b.worker_shards == ()
+    assert isinstance(b.engine, MultiDeviceEngine)
+
+
+def test_fit_predict_round_trip_and_residency_mirror(built):
+    engine = built.engine
+    engine.fit(*table(10, label="x"), dataset_id="rt")
+    holders = [w.device for w in built.worker_shards if w.holds("rt")]
+    assert len(holders) == 1  # mirrored before fit() returned
+    assert engine.predict("rt", [[0.0, 0.0]]).predictions == ["x"]
+    assert engine.feature_names("rt") is None
+    engine.evict("rt")
+    assert not any(w.has("rt") for w in built.worker_shards)
+
+
+def test_errors_cross_the_process_boundary_with_their_type(built):
+    with pytest.raises(DatasetNotFoundError):
+        built.engine.predict("never-fit", [[0.0, 0.0]])
+    with pytest.raises(InvalidInputError):
+        built.engine.fit(*table(10), dataset_id="b", backend="nope")
+    with pytest.raises(InvalidInputError):  # raised inside the worker
+        built.engine.fit([[1.0], [2.0]], ["a"], dataset_id="ragged")
+
+
+def test_concurrent_calls_from_many_threads(built):
+    engine = built.engine
+    for i in range(4):
+        engine.fit(*table(10, label=f"L{i}"), dataset_id=f"c{i}")
+    errors: list[Exception] = []
+
+    def hammer(i: int) -> None:
+        try:
+            for _ in range(20):
+                out = engine.predict(f"c{i % 4}", [[0.0, 0.0]])
+                assert out.predictions == [f"L{i % 4}"]
+        except Exception as e:  # noqa: BLE001
+            errors.append(e)
+
+    threads = [threading.Thread(target=hammer, args=(i,)) for i in range(8)]
+    for t in threads:
+        t.start()
+    for t in threads:
+        t.join(30)
+    assert errors == []
+    # Contexts spread over both workers (placement saw both as idle).
+    assert all(w.stats_nowait().n_cached_contexts >= 1 for w in built.worker_shards)
+
+
+def test_stats_and_reports(built):
+    built.engine.fit(*table(10), dataset_id="s")
+    stats = built.engine.stats()
+    assert stats.n_cached_contexts >= 1
+    reports = [w.report() for w in built.worker_shards]
+    assert all(r["gpu_memory"] is None for r in reports)  # "cuda:N" fake labels
+    assert all(r["spill"] is not None for r in reports)
+
+
+def test_a_dead_worker_fails_calls_and_is_skipped():
+    b = build_engine(_settings())
+    try:
+        engine = b.engine
+        victim, survivor = b.worker_shards
+        victim._proc.kill()
+        victim._proc.join(10)
+        for _ in range(200):  # the reader thread notices the closed pipe
+            if not victim.alive:
+                break
+            threading.Event().wait(0.01)
+        assert not victim.alive and "exited" in (victim.death_reason or "")
+        with pytest.raises(WorkerDiedError):
+            victim.fit(*table(5), dataset_id="x")
+        # The router places new work on the survivor only.
+        for i in range(3):
+            engine.fit(*table(5), dataset_id=f"after{i}")
+        assert all(survivor.holds(f"after{i}") for i in range(3))
+        assert engine.predict("after0", [[0.0, 0.0]]).predictions == ["a"]
+    finally:
+        b.close()
+
+
+def test_startup_failure_is_raised_in_the_serving_process(tmp_path):
+    # The worker can't create its spill subdirectory under a regular file.
+    blocker = tmp_path / "not-a-dir"
+    blocker.write_text("")
+    with pytest.raises(OSError):
+        build_engine(_settings(blocker))

From 31561eb1bb91f614f87d520738e7207c46fb54ca Mon Sep 17 00:00:00 2001
From: Amrit Krishnan <amrit110@gmail.com>
Date: Sat, 17 Oct 2026 19:40:41 -0400
Subject: [PATCH 15/39] feat(memory): TabPFN-3.5 requirement calibration
 (A100); width-dependent resident extrapolation

Measured with bench_memory_pressure --search --max-slowdown 2 (a fresh
process per shape): TabPFN, like TabICL, sizes its working memory from
what is free -- 20,000 x 100 peaks at 7.3 GiB unconstrained but runs
within 2x its time with 1.1 GiB free -- so admission now bounds resident
+ working floor for it too, instead of TabICL's peak formula (which held
TabPFN to 5k rows at 100 features on an L4).

TabPFN's per-row cache grows with features, so ResidentRateModel takes
width_dependent=True for it: the rate comes from the narrowest measured
width >= the query (short fits included, conservatively) and wider
shapes than any measurement get no estimate.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../baselines/minfree-tabpfn-kv-a100.jsonl    | 14 +++
 benchmarks/gke/serve_workload.sh              | 30 ++++++
 .../calibration_tabpfn_a100_requirement.py    | 91 +++++++++++++++++++
 src/tabctx/memory/requirement.py              | 56 ++++++++----
 src/tabctx/serve/factory.py                   | 10 +-
 tests/unit/test_requirement.py                | 50 ++++++++--
 6 files changed, 221 insertions(+), 30 deletions(-)
 create mode 100644 benchmarks/baselines/minfree-tabpfn-kv-a100.jsonl
 create mode 100755 benchmarks/gke/serve_workload.sh
 create mode 100644 src/tabctx/memory/calibration_tabpfn_a100_requirement.py

diff --git a/benchmarks/baselines/minfree-tabpfn-kv-a100.jsonl b/benchmarks/baselines/minfree-tabpfn-kv-a100.jsonl
new file mode 100644
index 0000000..562e1fc
--- /dev/null
+++ b/benchmarks/baselines/minfree-tabpfn-kv-a100.jsonl
@@ -0,0 +1,14 @@
+{"rows": 1000, "features": 10, "status": "ok", "resident_gib": 0.05, "unconstrained_peak_gib": 0.58, "unconstrained_fit_s": 0.172, "min_free_gib": 0.92, "min_working_gib": 0.87, "fit_s_at_min": 0.152, "max_slowdown": 2.0}
+{"rows": 10000, "features": 10, "status": "ok", "resident_gib": 0.31, "unconstrained_peak_gib": 2.26, "unconstrained_fit_s": 1.59, "min_free_gib": 2.91, "min_working_gib": 2.6, "fit_s_at_min": 1.587, "max_slowdown": 2.0}
+{"rows": 50000, "features": 10, "status": "ok", "resident_gib": 1.46, "unconstrained_peak_gib": 4.83, "unconstrained_fit_s": 16.42, "min_free_gib": 4.26, "min_working_gib": 2.8, "fit_s_at_min": 18.03, "max_slowdown": 2.0}
+{"rows": 100000, "features": 10, "status": "ok", "resident_gib": 2.89, "unconstrained_peak_gib": 9.64, "unconstrained_fit_s": 55.163, "min_free_gib": 8.54, "min_working_gib": 5.65, "fit_s_at_min": 58.662, "max_slowdown": 2.0}
+{"rows": 2000, "features": 100, "status": "ok", "resident_gib": 0.22, "unconstrained_peak_gib": 6.67, "unconstrained_fit_s": 0.845, "min_free_gib": 0.8, "min_working_gib": 0.58, "fit_s_at_min": 1.068, "max_slowdown": 2.0}
+{"rows": 5000, "features": 100, "status": "ok", "resident_gib": 0.34, "unconstrained_peak_gib": 6.69, "unconstrained_fit_s": 1.913, "min_free_gib": 1.21, "min_working_gib": 0.87, "fit_s_at_min": 2.707, "max_slowdown": 2.0}
+{"rows": 10000, "features": 100, "status": "ok", "resident_gib": 0.51, "unconstrained_peak_gib": 6.98, "unconstrained_fit_s": 4.115, "min_free_gib": 1.09, "min_working_gib": 0.58, "fit_s_at_min": 5.288, "max_slowdown": 2.0}
+{"rows": 20000, "features": 100, "status": "ok", "resident_gib": 0.8, "unconstrained_peak_gib": 7.3, "unconstrained_fit_s": 9.669, "min_free_gib": 1.09, "min_working_gib": 0.29, "fit_s_at_min": 11.485, "max_slowdown": 2.0}
+{"rows": 1000, "features": 500, "status": "ok", "resident_gib": 0.81, "unconstrained_peak_gib": 7.0, "unconstrained_fit_s": 2.11, "min_free_gib": 3.09, "min_working_gib": 2.28, "fit_s_at_min": 3.492, "max_slowdown": 2.0}
+{"rows": 2000, "features": 500, "status": "ok", "resident_gib": 0.88, "unconstrained_peak_gib": 7.16, "unconstrained_fit_s": 3.315, "min_free_gib": 3.16, "min_working_gib": 2.28, "fit_s_at_min": 5.236, "max_slowdown": 2.0}
+{"rows": 5000, "features": 500, "status": "ok", "resident_gib": 1.1, "unconstrained_peak_gib": 7.38, "unconstrained_fit_s": 8.321, "min_free_gib": 3.36, "min_working_gib": 2.26, "fit_s_at_min": 10.432, "max_slowdown": 2.0}
+{"rows": 500, "features": 1000, "status": "ok", "resident_gib": 1.18, "unconstrained_peak_gib": 7.31, "unconstrained_fit_s": 2.42, "min_free_gib": 5.7, "min_working_gib": 4.52, "fit_s_at_min": 3.863, "max_slowdown": 2.0}
+{"rows": 1000, "features": 1000, "status": "ok", "resident_gib": 1.22, "unconstrained_peak_gib": 7.42, "unconstrained_fit_s": 3.302, "min_free_gib": 5.73, "min_working_gib": 4.51, "fit_s_at_min": 4.81, "max_slowdown": 2.0}
+{"rows": 2000, "features": 1000, "status": "ok", "resident_gib": 1.32, "unconstrained_peak_gib": 7.64, "unconstrained_fit_s": 5.257, "min_free_gib": 5.82, "min_working_gib": 4.5, "fit_s_at_min": 6.9, "max_slowdown": 2.0}
diff --git a/benchmarks/gke/serve_workload.sh b/benchmarks/gke/serve_workload.sh
new file mode 100755
index 0000000..cae6b23
--- /dev/null
+++ b/benchmarks/gke/serve_workload.sh
@@ -0,0 +1,30 @@
+#!/usr/bin/env bash
+# Serve tabctx with Ray Serve inside the devbox pod (all of its GPUs, one
+# replica, as prod's tabctx worker group does) and run the mixed workload
+# against it. Run inside the pod (devbox.sh exec -- 'bash benchmarks/gke/serve_workload.sh ...').
+#
+#   serve_workload.sh OUT.json [USERS] [DURATION_S] [extra env assignments...]
+#   e.g. serve_workload.sh out/wl-procs.json 6 300 TABCTX_WORKER_PROCESSES=1
+set -euo pipefail
+OUT="$1"; USERS="${2:-6}"; DURATION="${3:-300}"; shift $(( $# < 3 ? $# : 3 ))
+for kv in "$@"; do export "$kv"; done
+export TABCTX_BACKEND="${TABCTX_BACKEND:-tabicl,tabpfn}"
+export TABCTX_SPILL_DIR="${TABCTX_SPILL_DIR:-/tmp/tabctx-spill}"
+# A num_gpus=0 actor gets CUDA_VISIBLE_DEVICES="" from Ray unless told not to
+# (prod sets the same; see inference-platform's tabctx worker group).
+export RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=1
+rm -rf "$TABCTX_SPILL_DIR"
+ray stop --force >/dev/null 2>&1 || true
+ray start --head --disable-usage-stats >/dev/null
+nohup serve run tabctx.serve.app:app > out/serve.log 2>&1 &
+SERVE_PID=$!
+trap 'kill $SERVE_PID 2>/dev/null || true; serve shutdown -y >/dev/null 2>&1 || true; ray stop --force >/dev/null 2>&1 || true' EXIT
+for _ in $(seq 1 180); do
+  if curl -sf http://127.0.0.1:8000/readyz > out/readyz-before.json; then break; fi
+  sleep 5
+done
+curl -sf http://127.0.0.1:8000/readyz > /dev/null || { tail -50 out/serve.log; exit 1; }
+python -c "import json; d=json.load(open('out/readyz-before.json')); print('ready:', d['device'], 'workers', len(d.get('devices') or []))"
+python benchmarks/bench_workload.py --base-url http://127.0.0.1:8000 \
+  --users "$USERS" --duration "$DURATION" --save "$OUT"
+curl -sf http://127.0.0.1:8000/readyz > "${OUT%.json}-readyz.json" || true
diff --git a/src/tabctx/memory/calibration_tabpfn_a100_requirement.py b/src/tabctx/memory/calibration_tabpfn_a100_requirement.py
new file mode 100644
index 0000000..5801c6d
--- /dev/null
+++ b/src/tabctx/memory/calibration_tabpfn_a100_requirement.py
@@ -0,0 +1,91 @@
+"""GENERATED by benchmarks/gen_requirement_grid.py -- do not hand-edit.
+
+Measured requirements (A100-SXM4-40GB, tabpfn 9.1.0 (TabPFN-3.5), fit_with_cache, max slowdown 2x): for each shape, the least free GPU
+memory with which a fit succeeded (resident context + working floor), from
+benchmarks/bench_memory_pressure.py --search. See memory/requirement.py.
+"""
+
+from __future__ import annotations
+
+from tabctx.memory.adaptive import Observation
+from tabctx.memory.requirement import FloorPoint
+
+A100_40GB_TABPFN_KV_FLOOR_POINTS: tuple[FloorPoint, ...] = (
+    FloorPoint(
+        n_train=1000, n_features=10, working_bytes=934155386, resident_bytes=53687091
+    ),  # resident 0.05 GiB, unconstrained peak 0.58 GiB, fit 0.172s -> 0.152s at the floor
+    FloorPoint(
+        n_train=10000, n_features=10, working_bytes=2791728742, resident_bytes=332859965
+    ),  # resident 0.31 GiB, unconstrained peak 2.26 GiB, fit 1.59s -> 1.587s at the floor
+    FloorPoint(
+        n_train=50000,
+        n_features=10,
+        working_bytes=3006477107,
+        resident_bytes=1567663063,
+    ),  # resident 1.46 GiB, unconstrained peak 4.83 GiB, fit 16.42s -> 18.03s at the floor
+    FloorPoint(
+        n_train=100000,
+        n_features=10,
+        working_bytes=6066641305,
+        resident_bytes=3103113871,
+    ),  # resident 2.89 GiB, unconstrained peak 9.64 GiB, fit 55.163s -> 58.662s at the floor
+    FloorPoint(
+        n_train=2000, n_features=100, working_bytes=622770257, resident_bytes=236223201
+    ),  # resident 0.22 GiB, unconstrained peak 6.67 GiB, fit 0.845s -> 1.068s at the floor
+    FloorPoint(
+        n_train=5000, n_features=100, working_bytes=934155386, resident_bytes=365072220
+    ),  # resident 0.34 GiB, unconstrained peak 6.69 GiB, fit 1.913s -> 2.707s at the floor
+    FloorPoint(
+        n_train=10000, n_features=100, working_bytes=622770257, resident_bytes=547608330
+    ),  # resident 0.51 GiB, unconstrained peak 6.98 GiB, fit 4.115s -> 5.288s at the floor
+    FloorPoint(
+        n_train=20000, n_features=100, working_bytes=311385128, resident_bytes=858993459
+    ),  # resident 0.8 GiB, unconstrained peak 7.3 GiB, fit 9.669s -> 11.485s at the floor
+    FloorPoint(
+        n_train=1000, n_features=500, working_bytes=2448131358, resident_bytes=869730877
+    ),  # resident 0.81 GiB, unconstrained peak 7.0 GiB, fit 2.11s -> 3.492s at the floor
+    FloorPoint(
+        n_train=2000, n_features=500, working_bytes=2448131358, resident_bytes=944892805
+    ),  # resident 0.88 GiB, unconstrained peak 7.16 GiB, fit 3.315s -> 5.236s at the floor
+    FloorPoint(
+        n_train=5000,
+        n_features=500,
+        working_bytes=2426656522,
+        resident_bytes=1181116006,
+    ),  # resident 1.1 GiB, unconstrained peak 7.38 GiB, fit 8.321s -> 10.432s at the floor
+    FloorPoint(
+        n_train=500,
+        n_features=1000,
+        working_bytes=4853313044,
+        resident_bytes=1267015352,
+    ),  # resident 1.18 GiB, unconstrained peak 7.31 GiB, fit 2.42s -> 3.863s at the floor
+    FloorPoint(
+        n_train=1000,
+        n_features=1000,
+        working_bytes=4842575626,
+        resident_bytes=1309965025,
+    ),  # resident 1.22 GiB, unconstrained peak 7.42 GiB, fit 3.302s -> 4.81s at the floor
+    FloorPoint(
+        n_train=2000,
+        n_features=1000,
+        working_bytes=4831838208,
+        resident_bytes=1417339207,
+    ),  # resident 1.32 GiB, unconstrained peak 7.64 GiB, fit 5.257s -> 6.9s at the floor
+)
+
+A100_40GB_TABPFN_KV_REQUIREMENT_GRID: tuple[Observation, ...] = (
+    Observation(n_train=1000, n_features=10, real_bytes=987842478),
+    Observation(n_train=10000, n_features=10, real_bytes=3124588707),
+    Observation(n_train=50000, n_features=10, real_bytes=4574140170),
+    Observation(n_train=100000, n_features=10, real_bytes=9169755176),
+    Observation(n_train=2000, n_features=100, real_bytes=858993459),
+    Observation(n_train=5000, n_features=100, real_bytes=1299227607),
+    Observation(n_train=10000, n_features=100, real_bytes=1170378588),
+    Observation(n_train=20000, n_features=100, real_bytes=1170378588),
+    Observation(n_train=1000, n_features=500, real_bytes=3317862236),
+    Observation(n_train=2000, n_features=500, real_bytes=3393024163),
+    Observation(n_train=5000, n_features=500, real_bytes=3607772528),
+    Observation(n_train=500, n_features=1000, real_bytes=6120328396),
+    Observation(n_train=1000, n_features=1000, real_bytes=6152540651),
+    Observation(n_train=2000, n_features=1000, real_bytes=6249177415),
+)
diff --git a/src/tabctx/memory/requirement.py b/src/tabctx/memory/requirement.py
index bf22b03..6dad7a2 100644
--- a/src/tabctx/memory/requirement.py
+++ b/src/tabctx/memory/requirement.py
@@ -103,24 +103,37 @@ def bytes(self, n_train: int, n_features: int) -> int | None:
 class ResidentRateModel:
     """Extrapolates a kv context's resident size to longer tables.
 
-    A TabICL kv context is a column cache that depends on features but not
-    rows, plus an ICL cache with one key/value per training row whose width
-    doesn't depend on features. So for any measured fit (n0 rows, d0 >= d
-    features, resident r0), resident(n, d) <= r0 + rate * (n - n0) for
-    n >= n0, where rate is the steepest per-row slope among long fits
-    (>= slope_min_rows; short fits are dominated by the column cache) at
-    the narrowest measured width >= d: TabPFN's per-row cache grows with
-    features, so only a rate measured at least as wide bounds it, and the
-    narrowest such is the tightest.
-    The estimate is the smallest such bound. Residents are fit-time sizes
-    (before any post-fit compression -- what the fit needs while it runs),
-    the same quantity the calibration measured. This replica's own long
-    fits supersede the calibration slope once any exist."""
+    For any measured fit (n0 rows, d0 >= d features, resident r0),
+    resident(n, d) <= r0 + rate * (n - n0) for n >= n0; the estimate is
+    the smallest such bound. What bounds the per-row rate depends on how
+    the model's cache scales:
+
+    - Width-independent (TabICL): a kv context is a column cache that
+      depends on features but not rows, plus an ICL cache with one
+      key/value per training row whose size doesn't depend on features.
+      The rate is the steepest among long fits (>= slope_min_rows; short
+      fits are dominated by the column cache), preferring the narrowest
+      measured width >= d.
+    - Width-dependent (TabPFN, ``width_dependent=True``): the per-row
+      cache grows with features, so only fits at least as wide bound the
+      rate. It comes from the narrowest measured width >= d -- long fits
+      there if any, else short ones (whose fixed overhead overstates the
+      rate: conservative) -- and shapes wider than every measurement get
+      no estimate at all.
+
+    Residents are fit-time sizes (before any post-fit compression -- what
+    the fit needs while it runs), the same quantity the calibration
+    measured. This replica's own fits supersede the calibration's rate
+    once any qualify."""
 
     def __init__(
-        self, calibration: Iterable[FloorPoint] = (), slope_min_rows: int = 10_000
+        self,
+        calibration: Iterable[FloorPoint] = (),
+        slope_min_rows: int = 10_000,
+        width_dependent: bool = False,
     ) -> None:
         self._min_rows = slope_min_rows
+        self._width_dependent = width_dependent
         self._calibration = [
             (p.n_train, p.n_features, p.resident_bytes)
             for p in calibration
@@ -138,12 +151,15 @@ def record(self, n_train: int, n_features: int, resident_bytes: int) -> None:
     def _slope(
         self, points: list[tuple[int, int, int]], n_features: int
     ) -> float | None:
-        """Steepest per-row rate among long fits at the narrowest measured
-        width >= n_features: residents don't shrink as features grow, so
-        that rate bounds this width's (and TabPFN's grows with width, so a
-        wider one would overshoot). Falls back to every long fit when none
-        is that wide (TabICL's rate barely depends on features)."""
-        long_fits = [(d, res / n) for n, d, res in points if n >= self._min_rows]
+        rates = [(n, d, res / n) for n, d, res in points]
+        long_fits = [(d, r) for n, d, r in rates if n >= self._min_rows]
+        if self._width_dependent:
+            wide = [d for _, d, _ in rates if d >= n_features]
+            if not wide:
+                return None
+            at = min(wide)
+            long_at = [r for d, r in long_fits if d == at]
+            return max(long_at or [r for _, d, r in rates if d == at])
         if not long_fits:
             return None
         widths = [d for d, _ in long_fits if d >= n_features]
diff --git a/src/tabctx/serve/factory.py b/src/tabctx/serve/factory.py
index 271bbb4..4614894 100644
--- a/src/tabctx/serve/factory.py
+++ b/src/tabctx/serve/factory.py
@@ -324,6 +324,10 @@ def _preloaded_observations(
     ("tabicl", "kv"): ("calibration_tabicl_a100_requirement", "A100_40GB_TABICL_KV"),
     ("tabpfn", "kv"): ("calibration_tabpfn_a100_requirement", "A100_40GB_TABPFN_KV"),
 }
+# Models whose per-row resident cache grows with features (see
+# ResidentRateModel): TabPFN caches per (row, feature group); TabICL's ICL
+# cache is per row only.
+_WIDTH_DEPENDENT_RESIDENT = frozenset({"tabpfn"})
 
 
 def _requirement_calibration(
@@ -457,7 +461,11 @@ def build_estimator(
         preloaded=tuple(fit_grid) + tuple(requirements),
         preloaded_predict=predict_grid,
         working_floor=working_floor,
-        resident_rate=ResidentRateModel(floor_points) if working_floor else None,
+        resident_rate=ResidentRateModel(
+            floor_points, width_dependent=kind in _WIDTH_DEPENDENT_RESIDENT
+        )
+        if working_floor
+        else None,
     )
 
 
diff --git a/tests/unit/test_requirement.py b/tests/unit/test_requirement.py
index 2428cf3..4c10b7e 100644
--- a/tests/unit/test_requirement.py
+++ b/tests/unit/test_requirement.py
@@ -235,14 +235,46 @@ def test_floor_model_is_monotone_in_rows():
     assert m.bytes(3_000, 1_000) >= m.bytes(2_000, 1_000)
 
 
-def test_factory_wires_floor_and_resident_models_for_tabicl_kv_only():
+def test_factory_wires_floor_and_resident_models_for_calibrated_kv_models():
     from tabctx.serve.factory import ServeSettings, build_estimator
 
-    kv = build_estimator(ServeSettings(backends=("tabicl",)), "tabicl", 40 * GiB)
-    assert kv._working_floor is not None and kv._resident_rate is not None
-    for settings, kind in (
-        (ServeSettings(backends=("tabicl",), kv_cache="repr"), "tabicl"),
-        (ServeSettings(backends=("tabpfn",)), "tabpfn"),
-    ):
-        est = build_estimator(settings, kind, 40 * GiB)
-        assert est._working_floor is None and est._resident_rate is None
+    for kind, width_dependent in (("tabicl", False), ("tabpfn", True)):
+        est = build_estimator(ServeSettings(backends=(kind,)), kind, 40 * GiB)
+        assert est._working_floor is not None and est._resident_rate is not None
+        assert est._resident_rate._width_dependent is width_dependent
+    repr_mode = build_estimator(
+        ServeSettings(backends=("tabicl",), kv_cache="repr"), "tabicl", 40 * GiB
+    )
+    assert repr_mode._working_floor is None and repr_mode._resident_rate is None
+
+
+class TestWidthDependentResidentRate:
+    """TabPFN-like caches: per-row size grows with features."""
+
+    CAL = (
+        FloorPoint(100_000, 10, GiB, resident_bytes=3 * GiB),
+        FloorPoint(20_000, 100, GiB, resident_bytes=1 * GiB),
+        FloorPoint(2_000, 1_000, GiB, resident_bytes=int(1.3 * GiB)),
+    )
+
+    def _model(self):
+        from tabctx.memory.requirement import ResidentRateModel
+
+        return ResidentRateModel(self.CAL, width_dependent=True)
+
+    def test_wider_than_any_measurement_has_no_estimate(self):
+        assert self._model().estimate(1_000, 2_000) is None
+
+    def test_short_fits_set_the_rate_when_no_long_fit_is_that_wide(self):
+        # Only 2000 rows measured at 1000 features: its whole resident
+        # counts as per-row (conservative), never the 10-feature rate.
+        rate = 1.3 * GiB / 2_000
+        assert self._model().estimate(20_000, 1_000) == pytest.approx(
+            1.3 * GiB + rate * 18_000, rel=1e-6
+        )
+
+    def test_narrow_queries_use_their_own_width(self):
+        rate = 3 * GiB / 100_000
+        assert self._model().estimate(150_000, 10) == pytest.approx(
+            3 * GiB + rate * 50_000, rel=1e-6
+        )

From d2316ecbe142b63b3df2f6fe22648d1b815155ca Mon Sep 17 00:00:00 2001
From: Amrit Krishnan <amrit110@gmail.com>
Date: Sat, 17 Oct 2026 19:47:16 -0400
Subject: [PATCH 16/39] fix(tabicl): one active-context cache per thread on the
 shared backbone

tabicl keeps the context a forward is using on the module (model_._cache)
and clears it after each call; with the backbone shared by every context,
a predict interleaved with a fit on the same GPU could read the fit's
cache (another tenant's answers) or None mid-forward (seen on an A100:
'NoneType' object has no attribute 'icl_cache'). The shared module now
gets a per-instance subclass whose _cache is thread-local. Today's
engines serialize every call on a device, so this was latent; it is the
prerequisite for interleaving predicts with fits (head-of-line, #39).

bench_overlap now checks predictions made during a fit match the idle
ones: 0 mismatches and 0 errors over 106 interleaved predicts (A100).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 benchmarks/bench_overlap.py                  | 10 +++-
 src/tabctx/backends/tabicl.py                | 37 ++++++++++++
 tests/unit/test_tabicl_thread_local_cache.py | 60 ++++++++++++++++++++
 3 files changed, 106 insertions(+), 1 deletion(-)
 create mode 100644 tests/unit/test_tabicl_thread_local_cache.py

diff --git a/benchmarks/bench_overlap.py b/benchmarks/bench_overlap.py
index 93bc3ca..aae2661 100644
--- a/benchmarks/bench_overlap.py
+++ b/benchmarks/bench_overlap.py
@@ -65,10 +65,17 @@ def main() -> int:
     small = backend.fit(Xc, yc, "classification")
     X_test, _ = make_table(args.test_rows, cf, "classification", seed=2)
 
+    reference = backend.predict(small, X_test, False).predictions
+    mismatches = [0]
+
     def predict_once() -> float:
         t0 = time.perf_counter()
-        backend.predict(small, X_test, False)
+        out = backend.predict(small, X_test, False)
         torch.cuda.synchronize(device)
+        # A predict sharing the backbone with a fit must still answer from
+        # its own context.
+        if out.predictions != reference:
+            mismatches[0] += 1
         return time.perf_counter() - t0
 
     for _ in range(3):
@@ -125,6 +132,7 @@ def predictor() -> None:
         "predict_during_fit": _pcts(during),
         "predict_errors": predict_errors[:5],
         "n_predict_errors": len(predict_errors),
+        "predict_mismatches": mismatches[0],
     }
     print(json.dumps(rec), flush=True)
     if args.out:
diff --git a/src/tabctx/backends/tabicl.py b/src/tabctx/backends/tabicl.py
index bd80fd0..16f0031 100644
--- a/src/tabctx/backends/tabicl.py
+++ b/src/tabctx/backends/tabicl.py
@@ -91,6 +91,42 @@ def _load_model(self: Any) -> None:
     cls._tabctx_restore_hook = True
 
 
+def make_cache_thread_local(module: Any) -> None:
+    """Give a shared TabICL backbone one active-context cache per thread.
+
+    tabicl's forward_with_cache keeps the context it is working with on
+    the MODULE (``model_._cache``), and every call clears it when done
+    (see _unpin_shared_module). With one backbone shared by every context,
+    two calls running at once -- a predict interleaved with a fit on the
+    same GPU -- would read each other's cache: a predict answered from
+    another tenant's context, or None mid-forward (seen on an A100 as
+    AttributeError: 'NoneType' object has no attribute 'icl_cache').
+    A per-instance subclass whose ``_cache`` is a thread-local property
+    gives each thread its own slot; nothing else about the module changes.
+    """
+    if getattr(type(module), "_tabctx_thread_local_cache", False):
+        return
+    local = threading.local()
+    base = type(module)
+
+    def _get(_self: Any) -> Any:
+        return getattr(local, "cache", None)
+
+    def _set(_self: Any, value: Any) -> None:
+        local.cache = value
+
+    module.__dict__.pop("_cache", None)
+    module.__class__ = type(
+        base.__name__,
+        (base,),
+        {
+            "_cache": property(_get, _set),
+            "_tabctx_thread_local_cache": True,
+            "__module__": base.__module__,
+        },
+    )
+
+
 class TabICLBackend:
     # The exact model id callers select via the API's `model` field --
     # named for the checkpoint actually loaded (tabicl-classifier-v2-*).
@@ -185,6 +221,7 @@ def _make_estimator(self, task: Task) -> Any:
     def _stash_backbone(self, task: Task, model: Any) -> None:
         if task in self._shared_backbones:
             return
+        make_cache_thread_local(model.model_)
         shared = {"model_": model.model_}
         # Keep save()/pickling of estimators working: __setstate__ only
         # rebuilds from weights when model_config_ exists (see tabicl
diff --git a/tests/unit/test_tabicl_thread_local_cache.py b/tests/unit/test_tabicl_thread_local_cache.py
new file mode 100644
index 0000000..0f9f42d
--- /dev/null
+++ b/tests/unit/test_tabicl_thread_local_cache.py
@@ -0,0 +1,60 @@
+"""The shared TabICL backbone keeps one active-context cache per thread
+(backends/tabicl.py make_cache_thread_local)."""
+
+import threading
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+from tabctx.backends.tabicl import make_cache_thread_local  # noqa: E402
+
+
+class Backbone(torch.nn.Module):
+    """Mimics tabicl's TabICL: _cache is a plain instance attribute set by
+    forward_with_cache and cleared by clear_cache."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.linear = torch.nn.Linear(2, 2)
+        self._cache = None
+
+    def forward(self, x):
+        return self.linear(x)
+
+    def clear_cache(self) -> None:
+        self._cache = None
+
+
+def test_each_thread_sees_its_own_cache():
+    m = Backbone()
+    make_cache_thread_local(m)
+    m._cache = "main"
+    seen: dict = {}
+    ready, go = threading.Event(), threading.Event()
+
+    def other() -> None:
+        assert m._cache is None  # never the main thread's
+        m._cache = "other"
+        ready.set()
+        go.wait(5)
+        seen["other"] = m._cache
+
+    t = threading.Thread(target=other)
+    t.start()
+    ready.wait(5)
+    m.clear_cache()  # the main thread's call ends: only its slot clears
+    go.set()
+    t.join(5)
+    assert seen["other"] == "other"
+    assert m._cache is None
+
+
+def test_module_still_works_and_is_idempotent():
+    m = Backbone()
+    make_cache_thread_local(m)
+    make_cache_thread_local(m)
+    assert isinstance(m, Backbone)
+    assert m(torch.zeros(1, 2)).shape == (1, 2)
+    assert "linear" in dict(m.named_modules())
+    assert len(list(m.parameters())) == 2

From 7c106bfc27823c3f5eda8f222f12d7e660693930 Mon Sep 17 00:00:00 2001
From: Amrit Krishnan <amrit110@gmail.com>
Date: Sat, 17 Oct 2026 19:52:29 -0400
Subject: [PATCH 17/39] feat(engine): light predicts run beside a fit on the
 same device (head-of-line)

Each device's engine used to hold its cache lock across every GPU call, so
a 100-row predict waited out any fit on that device -- up to 266 s on the
4xL4 rehearsal. Measured on an A100 (bench_overlap.py): a 100-row predict
during a 40k x 100 TabICL fit takes ~0.1 s p50 (1.5-2 s p95) and slows
the fit ~15%; TabPFN likewise (0.1 s p50, 0.86 s p95). 0 errors and 0
wrong answers across 208 interleaved predicts, with the thread-local
backbone cache (b715e91).

Scheduling per device:
- _fit_lock: fits, restores of spilled contexts and large predicts
  (> LIGHT_PREDICT_ROWS test rows) -- everything that grows residency.
- _predict_lock: one GPU predict at a time; a light predict on a resident
  context takes only this, so it overlaps a running fit.
- the cache lock now guards bookkeeping only (stats never wait on a fit).
- a predicting context is pinned: evict-ahead skips it, waiting for an
  unpin when nothing else can make room.
- a fit that hits OOM while sharing the device retries once alone.
- backends read allocator deltas under the predict lock
  (_cuda.QuiescedMeasurement), so a concurrent predict can't skew a
  fit's measured resident size.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/tabctx/backends/_cuda.py         |  22 +++
 src/tabctx/backends/tabicl.py        |  48 +++---
 src/tabctx/backends/tabpfn.py        |  26 +--
 src/tabctx/batching.py               |  15 +-
 src/tabctx/cache/manager.py          |  46 +++++-
 src/tabctx/engine.py                 | 227 ++++++++++++++++++---------
 src/tabctx/multidevice.py            |   8 +-
 src/tabctx/serve/app.py              |  12 +-
 src/tabctx/serve/workers.py          |   4 +-
 tests/unit/test_engine_scheduling.py | 203 ++++++++++++++++++++++++
 10 files changed, 483 insertions(+), 128 deletions(-)
 create mode 100644 tests/unit/test_engine_scheduling.py

diff --git a/src/tabctx/backends/_cuda.py b/src/tabctx/backends/_cuda.py
index 6dd7444..b414acb 100644
--- a/src/tabctx/backends/_cuda.py
+++ b/src/tabctx/backends/_cuda.py
@@ -37,3 +37,25 @@ def device_scope(device: str) -> contextlib.AbstractContextManager[Any]:
     import torch
 
     return torch.cuda.device(torch.device(device))
+
+
+class QuiescedMeasurement:
+    """Mixin for backends that measure a fit's footprint as a delta of
+    torch's allocator counters (memory_allocated before vs after).
+
+    The engine lets light predicts run beside a fit on the same device
+    (engine.py); a predict in flight at either reading would skew the
+    delta (in flight at the start: the context is understated -- the
+    unsafe direction). The engine lends its predict lock via
+    set_measurement_quiesce; readings taken under it see no predict's
+    temporaries (predicts free them before releasing the lock). Peaks
+    still include any predict that overlapped the fit: overstated,
+    which is the safe direction."""
+
+    _quiesce: Any = None
+
+    def set_measurement_quiesce(self, lock: Any) -> None:
+        self._quiesce = lock
+
+    def quiesced(self) -> contextlib.AbstractContextManager[Any]:
+        return self._quiesce if self._quiesce is not None else contextlib.nullcontext()
diff --git a/src/tabctx/backends/tabicl.py b/src/tabctx/backends/tabicl.py
index 16f0031..93f2bff 100644
--- a/src/tabctx/backends/tabicl.py
+++ b/src/tabctx/backends/tabicl.py
@@ -42,9 +42,9 @@
 needlessly throttled effective multi-tenant capacity to a fraction of what
 the hardware actually supports. Fixed here by measuring the real
 torch.cuda.memory_allocated() delta across fit() and reporting it via
-context_bytes_hint() -- safe because the engine holds its cache lock across
-the whole fit() call (see engine.py), so no concurrent GPU work can pollute
-the before/after delta.
+context_bytes_hint(). Predicts may run beside a fit on the same device, so
+the readings are taken under the engine's predict lock (see
+_cuda.QuiescedMeasurement): no predict's temporaries pollute the delta.
 """
 
 from __future__ import annotations
@@ -55,7 +55,12 @@
 
 import numpy as np
 
-from tabctx.backends._cuda import device_scope, ensure_initialized, is_cuda
+from tabctx.backends._cuda import (
+    QuiescedMeasurement,
+    device_scope,
+    ensure_initialized,
+    is_cuda,
+)
 from tabctx.backends.tabicl_kv import KV_PRECISIONS, KVPrecision, compress_icl_cache
 from tabctx.errors import BackendComputeError
 from tabctx.types import ArrayLike, PredictOutcome, Task
@@ -127,7 +132,7 @@ def _set(_self: Any, value: Any) -> None:
     )
 
 
-class TabICLBackend:
+class TabICLBackend(QuiescedMeasurement):
     # The exact model id callers select via the API's `model` field --
     # named for the checkpoint actually loaded (tabicl-classifier-v2-*).
     name = "tabicl-v2"
@@ -179,10 +184,10 @@ def __init__(
         # is tabicl's own sanctioned pattern (see
         # _unsupervised/unsupervised.py and _finetune/base.py, both of
         # which do exactly this with the comment "prevents redundant
-        # torch.load() calls"). Safe here because the engine serializes
-        # every backend call behind its cache lock -- the shared module's
-        # per-call mutable state (its _cache pointer, InferenceManager
-        # config) is never touched concurrently.
+        # torch.load() calls"). Calls on the shared module may overlap (a
+        # predict beside a fit, see engine.py): its active-context cache is
+        # per thread (make_cache_thread_local), and the InferenceManager
+        # config every call reapplies is identical for this backend.
         self._shared_backbones: dict[str, dict[str, Any]] = {}
         from tabicl import TabICLClassifier, TabICLRegressor
 
@@ -241,9 +246,11 @@ def fit(self, X: ArrayLike, y: ArrayLike, task: Task) -> Any:
 
         device = self._device
         measuring = is_cuda(device)
-        before_bytes = torch.cuda.memory_allocated(device) if measuring else 0
+        before_bytes = 0
         if measuring:
-            torch.cuda.reset_peak_memory_stats(device)
+            with self.quiesced():
+                before_bytes = torch.cuda.memory_allocated(device)
+                torch.cuda.reset_peak_memory_stats(device)
         self._last_context_bytes = None
         self._last_fit_peak_bytes = None
         self._last_fit_resident_bytes = None
@@ -253,11 +260,12 @@ def fit(self, X: ArrayLike, y: ArrayLike, task: Task) -> Any:
                 # The context as the fit built it, at full precision: what
                 # the fit needed resident, and so what admission must
                 # learn -- compression only shrinks it afterwards.
-                fit_resident = (
-                    torch.cuda.memory_allocated(device) - before_bytes
-                    if measuring
-                    else None
-                )
+                fit_resident = None
+                if measuring:
+                    with self.quiesced():
+                        fit_resident = (
+                            torch.cuda.memory_allocated(device) - before_bytes
+                        )
                 compress_icl_cache(model, self._kv_precision)
         except torch.cuda.OutOfMemoryError as e:
             raise BackendComputeError(f"CUDA OOM during fit(): {e}") from e
@@ -274,12 +282,12 @@ def fit(self, X: ArrayLike, y: ArrayLike, task: Task) -> Any:
         # - peak (fit_peak_bytes_hint): the transient high-water during fit.
         # None on CPU; the engine falls back to the formula estimator.
         if measuring:
-            after_bytes = torch.cuda.memory_allocated(device)
+            with self.quiesced():
+                after_bytes = torch.cuda.memory_allocated(device)
+                peak_bytes = torch.cuda.max_memory_allocated(device)
             self._last_context_bytes = max(0, after_bytes - before_bytes)
             self._last_fit_resident_bytes = max(0, fit_resident or 0)
-            self._last_fit_peak_bytes = max(
-                0, torch.cuda.max_memory_allocated(device) - before_bytes
-            )
+            self._last_fit_peak_bytes = max(0, peak_bytes - before_bytes)
         self._stash_backbone(task, model)
         return model
 
diff --git a/src/tabctx/backends/tabpfn.py b/src/tabctx/backends/tabpfn.py
index dfa86f3..8deadb4 100644
--- a/src/tabctx/backends/tabpfn.py
+++ b/src/tabctx/backends/tabpfn.py
@@ -38,7 +38,12 @@
 
 import numpy as np
 
-from tabctx.backends._cuda import device_scope, ensure_initialized, is_cuda
+from tabctx.backends._cuda import (
+    QuiescedMeasurement,
+    device_scope,
+    ensure_initialized,
+    is_cuda,
+)
 from tabctx.errors import BackendComputeError, InvalidInputError
 from tabctx.types import ArrayLike, PredictOutcome, Task
 
@@ -84,7 +89,7 @@ def persistent_load(self, pid: Any) -> Any:
         return self._resolve(pid)
 
 
-class TabPFNBackend:
+class TabPFNBackend(QuiescedMeasurement):
     def __init__(self, device: str | None = None, cache_mode: str = "kv") -> None:
         import torch
         from tabpfn.settings import settings
@@ -124,9 +129,11 @@ def fit(self, X: ArrayLike, y: ArrayLike, task: Task) -> Any:
         model = self._make_estimator(task)
         device = self._device
         measuring = is_cuda(device)
-        before_bytes = torch.cuda.memory_allocated(device) if measuring else 0
+        before_bytes = 0
         if measuring:
-            torch.cuda.reset_peak_memory_stats(device)
+            with self.quiesced():
+                before_bytes = torch.cuda.memory_allocated(device)
+                torch.cuda.reset_peak_memory_stats(device)
         try:
             with device_scope(device):
                 model.fit(np.asarray(X, dtype=float), np.asarray(y))
@@ -143,12 +150,11 @@ def fit(self, X: ArrayLike, y: ArrayLike, task: Task) -> Any:
         # Resident vs peak: two different consumers -- see
         # backends/tabicl.py fit() for the full rationale.
         if measuring:
-            self._last_context_bytes = max(
-                0, torch.cuda.memory_allocated(device) - before_bytes
-            )
-            self._last_fit_peak_bytes = max(
-                0, torch.cuda.max_memory_allocated(device) - before_bytes
-            )
+            with self.quiesced():
+                after_bytes = torch.cuda.memory_allocated(device)
+                peak_bytes = torch.cuda.max_memory_allocated(device)
+            self._last_context_bytes = max(0, after_bytes - before_bytes)
+            self._last_fit_peak_bytes = max(0, peak_bytes - before_bytes)
         else:
             self._last_context_bytes = None
             self._last_fit_peak_bytes = None
diff --git a/src/tabctx/batching.py b/src/tabctx/batching.py
index 2639df1..68f2a94 100644
--- a/src/tabctx/batching.py
+++ b/src/tabctx/batching.py
@@ -35,13 +35,14 @@
 first real-user deployment is what will measure it; until then this
 module's same-context coalescing is the batching layer.
 
-Safety argument (why this can't blow the memory budget): the engine
-serializes all GPU work behind its cache lock, and engine.predict()
-already chunks oversized test sets against the memory ceiling
-(chunking.py). Coalescing therefore never increases the number of
-concurrent GPU calls (still exactly one) -- it only reduces how many
-serialized calls the same amount of work needs. The admission ceiling
-derivation ("exactly one in-flight backend call") still holds.
+Safety argument (why this can't blow the memory budget): the engine runs
+at most one predict per device at a time (beside at most one fit, and only
+when the predict is light -- see engine.py), and engine.predict() already
+chunks oversized test sets against the memory ceiling (chunking.py).
+Coalescing therefore never increases the number of concurrent GPU calls --
+it only reduces how many serialized calls the same amount of work needs.
+A coalesced batch larger than LIGHT_PREDICT_ROWS simply waits for the
+device's fits, like any large predict.
 
 Failure isolation: one malformed request must not poison the others in
 its batch. If a *batched* call fails input validation, every member is
diff --git a/src/tabctx/cache/manager.py b/src/tabctx/cache/manager.py
index c2bf4f5..0ec56f8 100644
--- a/src/tabctx/cache/manager.py
+++ b/src/tabctx/cache/manager.py
@@ -19,8 +19,10 @@
 
 from __future__ import annotations
 
+import contextlib
 import threading
 import time
+from collections.abc import Iterator
 from dataclasses import dataclass, field
 from typing import Any
 
@@ -72,6 +74,12 @@ def __init__(
         # (has) must still see them: a multi-device router that saw
         # neither would answer 404 for a context that is merely moving.
         self._in_transit: set[str] = set()
+        # Contexts a call is using right now (see pinned()): never chosen as
+        # eviction victims, so a fit's evict-ahead can't spill a context out
+        # from under a predict running beside it -- the predict would keep
+        # the GPU memory alive while the books count it free.
+        self._pins: dict[str, int] = {}
+        self._unpinned = threading.Condition(self._lock)
 
     def get(self, dataset_id: str) -> CachedContext | None:
         with self._lock:
@@ -99,6 +107,40 @@ def get(self, dataset_id: str) -> CachedContext | None:
             finally:
                 self._in_transit.discard(dataset_id)
 
+    @contextlib.contextmanager
+    def pinned(self, dataset_id: str) -> Iterator[None]:
+        """Keep a context out of eviction's reach for the duration (explicit
+        evict() and re-fits still replace it: caller intent)."""
+        with self._lock:
+            self._pins[dataset_id] = self._pins.get(dataset_id, 0) + 1
+        try:
+            yield
+        finally:
+            with self._lock:
+                left = self._pins[dataset_id] - 1
+                if left:
+                    self._pins[dataset_id] = left
+                else:
+                    del self._pins[dataset_id]
+                    self._unpinned.notify_all()
+
+    def wait_for_unpin(self, timeout_s: float) -> bool:
+        """Wait until some pinned context is released. False if nothing is
+        pinned (nothing to wait for) or on timeout."""
+        with self._unpinned:
+            if not self._pins:
+                return False
+            before = sum(self._pins.values())
+            return self._unpinned.wait_for(
+                lambda: sum(self._pins.values()) < before, timeout_s
+            )
+
+    def peek(self, dataset_id: str) -> CachedContext | None:
+        """The resident context, or None -- never reloads a spilled one (a
+        restore needs device memory, which only an exclusive call may take)."""
+        with self._lock:
+            return self._entries.get(dataset_id)
+
     def touch(self, dataset_id: str) -> None:
         with self._lock:
             entry = self._entries.get(dataset_id)
@@ -116,7 +158,9 @@ def _remove(self, dataset_id: str) -> CachedContext | None:
     def _evict_victim(self) -> str | None:
         """Policy-chosen eviction, spilling when a tier is attached.
         Caller holds the lock."""
-        victim_id = self._policy.select_victim(list(self._entries.values()))
+        victim_id = self._policy.select_victim(
+            [e for e in self._entries.values() if e.dataset_id not in self._pins]
+        )
         if victim_id is None:
             return None
         if self._spill is None:
diff --git a/src/tabctx/engine.py b/src/tabctx/engine.py
index 6fd73e1..d98f902 100644
--- a/src/tabctx/engine.py
+++ b/src/tabctx/engine.py
@@ -7,6 +7,7 @@
 from __future__ import annotations
 
 import contextlib
+import threading
 import uuid
 
 from tabctx.backends.base import TabularICLBackend
@@ -22,6 +23,17 @@
 from tabctx.memory.host import HostMemoryGuard
 from tabctx.types import ArrayLike, EngineStats, PredictOutcome, Task
 
+# Predicts of at most this many test rows against a resident context run
+# beside a fit on the same device instead of waiting for it (head-of-line,
+# #39): measured on an A100, a 100-row predict took ~0.1 s p50 during a
+# 40k x 100 fit (vs the fit's full 11 s), slowing the fit ~15%. Larger
+# predicts, and predicts that must first restore a spilled context, need
+# device memory a running fit may be using, so they wait their turn.
+LIGHT_PREDICT_ROWS = 2048
+# How long a fit's evict-ahead waits for pinned (in-use) contexts to be
+# released before trying its luck with the headroom it has.
+UNPIN_WAIT_S = 30.0
+
 
 def _rect_shape(X: ArrayLike) -> tuple[int, int] | None:
     """(n_rows, n_cols) when X is a 2D array-with-shape (e.g. numpy) --
@@ -149,6 +161,20 @@ def __init__(
         self._estimators = estimators
         self._default_backend = default_backend
         self._cache = cache
+        # Device scheduling. _fit_lock serializes everything that grows
+        # residency on this device: fits, restores of spilled contexts and
+        # large predicts. _predict_lock runs one GPU predict at a time, so
+        # a light predict (see LIGHT_PREDICT_ROWS) overlaps a running fit
+        # rather than queueing behind it. The cache's own lock only guards
+        # bookkeeping. Lock order: _fit_lock, then _predict_lock.
+        self._fit_lock = threading.Lock()
+        # Re-entrant: a fit retrying alone holds it while the backend's
+        # measurement takes it again (see _cuda.QuiescedMeasurement).
+        self._predict_lock = threading.RLock()
+        for b in backends.values():
+            lend = getattr(b, "set_measurement_quiesce", None)
+            if callable(lend):
+                lend(self._predict_lock)
         # Process-wide host-memory admission, shared by every device's
         # engine (see memory/host.py); None = unguarded.
         self._host_guard = host_guard
@@ -262,75 +288,17 @@ def fit(
         estimator = self._estimators[backend_name]
 
         resolved_id = dataset_id or str(uuid.uuid4())
-        # Host memory first, outside the cache lock: a large fit waiting for
-        # a host slot must not block predicts on this device meanwhile.
-        with self._host_slot(chosen, n_train, n_features), self._cache.lock:
-            # Usage-aware admission (v0.9.0): what OOMs is the fit's
-            # transient PEAK on top of everything resident, so the
-            # estimated peak must fit the headroom given current cache
-            # usage -- checked under the cache lock so two concurrent
-            # fits can't both pass against the same snapshot.
-            #
-            # Evict-ahead-of-fit: when the cache is warm, cold contexts
-            # are evicted (spilled, when a spill tier is attached) until
-            # the fit's transient need fits -- without this, a filling
-            # replica would reject fits that are perfectly safe once a
-            # cold context steps aside. Proven necessary on real
-            # hardware: a fourth ~13GB-peak fit OOMed an A100 with three
-            # ~8GB contexts resident, exactly the state this drains.
+        # Host memory first: a large fit waiting for a host slot must not
+        # hold this device meanwhile.
+        with self._host_slot(chosen, n_train, n_features), self._fit_lock:
             estimated = estimator.estimate_bytes(n_train, 0, n_features)
-            # Feasibility FIRST, against an EMPTY cache's headroom: a fit
-            # that can never be admitted must be rejected before evicting
-            # anything -- otherwise one oversized request drains every
-            # tenant's context on its way to a 413 (observed live: a
-            # 200k x 60 request spilled 30 contexts and then failed).
-            max_headroom = self._headroom(estimator, 0)
-            if estimated > max_headroom:
-                raise AdmissionRejected(
-                    f"training shape ({n_train} rows x {n_features} "
-                    f"features) is estimated to need {estimated} bytes at "
-                    f"peak, exceeding this replica's "
-                    f"{max_headroom} byte maximum "
-                    "admission headroom (even with an empty cache)"
-                )
-            while True:
-                headroom = self._headroom(estimator, self._cache.stats().used_bytes)
-                if estimated <= headroom:
-                    break
-                # Feasible but blocked by warm contexts: evict-ahead.
-                # Terminates at the feasibility bound above; the None
-                # check is belt-and-braces against an estimator whose
-                # headroom doesn't recover to the empty-cache value.
-                if self._cache.evict_one() is None:
-                    break
-            payload = chosen.fit(X, y, task)
+            self._admit(estimator, estimated, n_train, n_features)
+            payload = self._fit_backend(chosen, X, y, task)
             est_bytes = chosen.context_bytes_hint(n_train, n_features)
             if est_bytes is None:
                 est_bytes = estimated
             else:
-                # Feed real measurements back so the PRE-FIT admission gate
-                # can use them (safely, as a bound for smaller/equal future
-                # shapes -- see memory/adaptive.py). The estimator decides
-                # what the fit needed: resident + working floor for models
-                # whose working memory adapts to what is free (their peak
-                # only says how much was free), else the transient peak.
-                peak_hint = getattr(chosen, "fit_peak_bytes_hint", None)
-                peak_bytes = peak_hint() if callable(peak_hint) else None
-                # What the fit needed resident while it ran: backends that
-                # compress the context afterwards (TabICL's int8 kv cache)
-                # report the pre-compression size; est_bytes, the compressed
-                # size, stays the cache-accounting quantity.
-                fit_hint = getattr(chosen, "fit_resident_bytes_hint", None)
-                fit_resident = fit_hint() if callable(fit_hint) else None
-                record_fit = getattr(estimator, "record_fit", None)
-                if callable(record_fit):
-                    record_fit(
-                        n_train, n_features, fit_resident or est_bytes, peak_bytes
-                    )
-                else:
-                    estimator.record_observation(
-                        n_train, n_features, peak_bytes if peak_bytes else est_bytes
-                    )
+                self._learn(chosen, estimator, n_train, n_features, est_bytes)
             context = CachedContext(
                 dataset_id=resolved_id,
                 backend_name=backend_name,
@@ -344,6 +312,92 @@ def fit(
             self._cache.put(context)
         return resolved_id
 
+    def _admit(
+        self,
+        estimator: MemoryEstimator,
+        estimated: int,
+        n_train: int,
+        n_features: int,
+    ) -> None:
+        """Make room for a fit's estimated need, or reject it. Caller holds
+        _fit_lock, so no other residency growth races this check.
+
+        Usage-aware admission (v0.9.0): what OOMs is the fit's need on top
+        of everything resident, so it must fit the headroom given current
+        cache usage. Evict-ahead-of-fit: when the cache is warm, cold
+        contexts are evicted (spilled, when a spill tier is attached) until
+        the need fits -- without this, a filling replica would reject fits
+        that are perfectly safe once a cold context steps aside (proven on
+        real hardware: a fourth ~13GB-peak fit OOMed an A100 with three
+        ~8GB contexts resident). Contexts a predict is using are pinned and
+        skipped; if only pinned ones stand in the way, wait for them."""
+        # Feasibility FIRST, against an EMPTY cache's headroom: a fit that
+        # can never be admitted must be rejected before evicting anything --
+        # otherwise one oversized request drains every tenant's context on
+        # its way to a 413 (observed live: a 200k x 60 request spilled 30
+        # contexts and then failed).
+        max_headroom = self._headroom(estimator, 0)
+        if estimated > max_headroom:
+            raise AdmissionRejected(
+                f"training shape ({n_train} rows x {n_features} "
+                f"features) is estimated to need {estimated} bytes at "
+                f"peak, exceeding this replica's "
+                f"{max_headroom} byte maximum "
+                "admission headroom (even with an empty cache)"
+            )
+        with self._cache.lock:
+            while True:
+                headroom = self._headroom(estimator, self._cache.stats().used_bytes)
+                if estimated <= headroom:
+                    return
+                if self._cache.evict_one() is not None:
+                    continue
+                # Only pinned contexts left (or none at all): wait for a
+                # predict to finish, else go ahead with what there is.
+                if not self._cache.wait_for_unpin(UNPIN_WAIT_S):
+                    return
+
+    def _fit_backend(
+        self, backend: TabularICLBackend, X: ArrayLike, y: ArrayLike, task: Task
+    ) -> object:
+        """Fit, retrying once with predicts held off if it ran out of
+        device memory while sharing the device with one."""
+        try:
+            return backend.fit(X, y, task)
+        except BackendComputeError:
+            with self._predict_lock:
+                return backend.fit(X, y, task)
+
+    def _learn(
+        self,
+        backend: TabularICLBackend,
+        estimator: MemoryEstimator,
+        n_train: int,
+        n_features: int,
+        est_bytes: int,
+    ) -> None:
+        """Feed a fit's real measurements back so the PRE-FIT admission gate
+        can use them (safely, as a bound for smaller/equal future shapes --
+        see memory/adaptive.py). The estimator decides what the fit needed:
+        resident + working floor for models whose working memory adapts to
+        what is free (their peak only says how much was free), else the
+        transient peak."""
+        peak_hint = getattr(backend, "fit_peak_bytes_hint", None)
+        peak_bytes = peak_hint() if callable(peak_hint) else None
+        # What the fit needed resident while it ran: backends that compress
+        # the context afterwards (TabICL's int8 kv cache) report the
+        # pre-compression size; est_bytes, the compressed size, stays the
+        # cache-accounting quantity.
+        fit_hint = getattr(backend, "fit_resident_bytes_hint", None)
+        fit_resident = fit_hint() if callable(fit_hint) else None
+        record_fit = getattr(estimator, "record_fit", None)
+        if callable(record_fit):
+            record_fit(n_train, n_features, fit_resident or est_bytes, peak_bytes)
+        else:
+            estimator.record_observation(
+                n_train, n_features, peak_bytes if peak_bytes else est_bytes
+            )
+
     def _host_slot(self, backend: TabularICLBackend, n_rows: int, n_features: int):
         """A process-wide large-call slot when this call may offload to host
         memory (the backend's may_offload, else the guard's cell threshold);
@@ -362,9 +416,11 @@ def feature_names(self, dataset_id: str) -> list[str] | None:
         """Column names recorded at fit() for a cached context, or None
         when the context is unknown or was fit without names. Reloads a
         spilled context as a side effect, exactly as predict() would."""
-        with self._cache.lock:
-            context = self._cache.get(dataset_id)
-            return context.feature_names if context is not None else None
+        context = self._cache.peek(dataset_id)
+        if context is None:
+            with self._fit_lock:  # a restore takes device memory
+                context = self._cache.get(dataset_id)
+        return context.feature_names if context is not None else None
 
     def predict(
         self, dataset_id: str, X_test: ArrayLike, return_proba: bool = False
@@ -374,21 +430,36 @@ def predict(
         (see chunking.py) against the memory ceiling -- this is the one
         thing that must never let a single request take down the replica,
         which the naive one-shot wrapper this library replaces did not do.
+
+        A light predict (a resident context, at most LIGHT_PREDICT_ROWS test
+        rows) runs beside a fit on this device; anything else waits for the
+        device's fits.
         """
-        context = self._cache.get(dataset_id)
-        if context is None:
-            raise DatasetNotFoundError(
-                f"no cached context for dataset_id={dataset_id!r} -- it was "
-                "never fit, was evicted, or was lost to a replica restart "
-                "(the cache has no durability across process restarts)"
-            )
+        with self._cache.pinned(dataset_id):
+            context = self._cache.peek(dataset_id)
+            if context is not None and len(X_test) <= LIGHT_PREDICT_ROWS:
+                return self._predict_on(context, X_test, return_proba)
+        with self._fit_lock, self._cache.pinned(dataset_id):
+            context = self._cache.get(dataset_id)  # restores a spilled one
+            if context is None:
+                raise DatasetNotFoundError(
+                    f"no cached context for dataset_id={dataset_id!r} -- it was "
+                    "never fit, was evicted, or was lost to a replica restart "
+                    "(the cache has no durability across process restarts)"
+                )
+            return self._predict_on(context, X_test, return_proba)
+
+    def _predict_on(
+        self, context: CachedContext, X_test: ArrayLike, return_proba: bool
+    ) -> PredictOutcome:
+        """Run a predict against a context the caller has pinned."""
         _validate_predict_input(X_test, context.n_features)
         # Dispatch to whichever backend fit this context (multi-backend
         # deployments serve several models over one cache).
         backend = self._backends.get(context.backend_name)
         if backend is None:
             raise DatasetNotFoundError(
-                f"context for {dataset_id!r} was fit by backend "
+                f"context for {context.dataset_id!r} was fit by backend "
                 f"{context.backend_name!r}, which this deployment no longer "
                 "serves -- re-fit with one of "
                 f"{sorted(self._backends)}"
@@ -403,7 +474,7 @@ def predict(
             estimator,
             context.n_train,
             context.n_features,
-            self._headroom(estimator, self._cache.stats().used_bytes),
+            self._headroom(estimator, self._cache.stats_nowait().used_bytes),
         )
         chunks = split_rows(X_test, chunk_rows) if chunk_rows < n_test else [X_test]
 
@@ -413,9 +484,9 @@ def predict(
         largest_chunk = max(len(c) for c in chunks)
         with (
             self._host_slot(backend, largest_chunk, context.n_features),
-            self._cache.lock,
+            self._predict_lock,
         ):
-            self._cache.touch(dataset_id)
+            self._cache.touch(context.dataset_id)
             for chunk in chunks:
                 for outcome in _predict_splitting_on_oom(
                     backend, context.payload, chunk, return_proba
diff --git a/src/tabctx/multidevice.py b/src/tabctx/multidevice.py
index bd56b94..9b1f382 100644
--- a/src/tabctx/multidevice.py
+++ b/src/tabctx/multidevice.py
@@ -13,7 +13,7 @@
     1. Devices where the table's estimated peak fits WITHOUT evicting
        anything are preferred; among them, the least active one wins
        (fewest calls in flight, then fewest started recently), ties to
-       the most headroom. Each device serializes its own GPU work, so
+       the most headroom. Each device runs one fit at a time, so
        spreading active users across devices is what lets them run
        concurrently.
     2. If no device can take it without eviction (a large table, or a
@@ -27,9 +27,9 @@
   device's residency and spill index (so a context spilled before a
   restart is found on, and restored to, its own device).
 
-Concurrency. Each device's engine holds its cache lock across its whole
-GPU call, so this router never takes another device's lock to make a
-decision: placement and lookup read lock-free snapshots (used bytes,
+Concurrency. Each device's engine schedules its own GPU work (one fit at a
+time, light predicts beside it; see engine.py), and this router never
+waits on a device to make a decision: placement and lookup read lock-free snapshots (used bytes,
 residency, spill index). Work on different devices therefore runs
 concurrently, fits included. Per-dataset_id locks serialize fits and
 evicts of the SAME id, so two concurrent fits of one id end with a single
diff --git a/src/tabctx/serve/app.py b/src/tabctx/serve/app.py
index bf5a5da..476f219 100644
--- a/src/tabctx/serve/app.py
+++ b/src/tabctx/serve/app.py
@@ -179,12 +179,12 @@ def _replica_tag() -> str | None:
 
 
 @serve.deployment(
-    # 8 concurrent requests per replica (was 2): safe because the engine
-    # serializes all GPU work behind its cache lock regardless -- extra
-    # in-flight requests wait on the lock (or coalesce, see batching.py)
-    # rather than running concurrent GPU calls, so the memory ceiling's
-    # one-in-flight-call assumption still holds. More in-flight requests
-    # = more same-context coalescing opportunity.
+    # 8 concurrent requests per replica (was 2): safe because each device's
+    # engine schedules its GPU work itself (one fit, plus one light predict
+    # beside it; see engine.py) -- extra in-flight requests queue there (or
+    # coalesce, see batching.py) rather than running unbounded concurrent
+    # GPU calls. More in-flight requests = more same-context coalescing
+    # opportunity.
     max_ongoing_requests=8,
     max_queued_requests=16,
     health_check_period_s=30,
diff --git a/src/tabctx/serve/workers.py b/src/tabctx/serve/workers.py
index 412fe64..1add8e6 100644
--- a/src/tabctx/serve/workers.py
+++ b/src/tabctx/serve/workers.py
@@ -62,8 +62,8 @@
         "report",
     }
 )
-# Concurrent calls per worker: its engine serializes GPU work itself
-# (cache lock); these only need to queue and answer cheap calls promptly.
+# Concurrent calls per worker: its engine schedules GPU work itself (see
+# engine.py); these only need to queue and answer cheap calls promptly.
 _WORKER_THREADS = 16
 
 
diff --git a/tests/unit/test_engine_scheduling.py b/tests/unit/test_engine_scheduling.py
new file mode 100644
index 0000000..a746df8
--- /dev/null
+++ b/tests/unit/test_engine_scheduling.py
@@ -0,0 +1,203 @@
+"""Per-device scheduling in TabctxEngine: light predicts overlap a running
+fit; large predicts, restores and fits serialize; evict-ahead never takes a
+context a predict is using. Interleavings are forced with gates, never
+timing assertions."""
+
+import threading
+
+import pytest
+
+from tabctx.backends.fake import FakeBackend
+from tabctx.cache.manager import ContextCacheManager
+from tabctx.cache.spill import DiskSpillStore
+from tabctx.engine import LIGHT_PREDICT_ROWS, TabctxEngine
+from tabctx.errors import BackendComputeError
+
+BYTES_PER_CELL = 10
+CAPACITY = 10_000
+
+
+class CellEstimator:
+    """Need = resident = cells x BYTES_PER_CELL; headroom = what the cache
+    doesn't hold."""
+
+    def estimate_bytes(self, n_train, n_test, n_features):
+        return (n_train + n_test) * n_features * BYTES_PER_CELL
+
+    def admit(self, n_train, n_test, n_features):
+        return self.estimate_bytes(n_train, n_test, n_features) <= CAPACITY
+
+    def ceiling_bytes(self):
+        return CAPACITY
+
+    def admission_headroom_bytes(self, used_bytes):
+        return max(0, CAPACITY - used_bytes)
+
+    def confidence(self):
+        return "test"
+
+    def record_observation(self, n_train, n_features, real_bytes):
+        pass
+
+
+class GatedFake(FakeBackend):
+    def __init__(self) -> None:
+        super().__init__()
+        self.fit_gate: threading.Event | None = None
+        self.predict_gate: threading.Event | None = None
+        self.fit_entered = threading.Event()
+        self.predict_entered = threading.Event()
+        self.fail_next_fits = 0
+        self._last = None
+
+    def fit(self, X, y, task):
+        self._last = len(X) * len(X[0]) * BYTES_PER_CELL
+        if self.fit_gate is not None:
+            self.fit_entered.set()
+            assert self.fit_gate.wait(5), "fit gate never released"
+        if self.fail_next_fits:
+            self.fail_next_fits -= 1
+            raise BackendComputeError("CUDA OOM during fit()")
+        return super().fit(X, y, task)
+
+    def predict(self, payload, X_test, return_proba=False):
+        if self.predict_gate is not None:
+            self.predict_entered.set()
+            assert self.predict_gate.wait(5), "predict gate never released"
+        return super().predict(payload, X_test, return_proba)
+
+    def context_bytes_hint(self, n_train, n_features):
+        return self._last
+
+
+def build(tmp_path=None):
+    backend = GatedFake()
+    spill = DiskSpillStore(tmp_path) if tmp_path else None
+    engine = TabctxEngine(
+        backend=backend,
+        estimator=CellEstimator(),
+        cache=ContextCacheManager(CAPACITY, spill_store=spill),
+    )
+    return engine, backend
+
+
+def table(rows, label="a"):
+    return [[float(i)] for i in range(rows)], [label] * rows
+
+
+def start(fn, *args, **kw):
+    box: dict = {}
+
+    def run():
+        try:
+            box["result"] = fn(*args, **kw)
+        except Exception as e:  # noqa: BLE001
+            box["error"] = e
+
+    t = threading.Thread(target=run, daemon=True)
+    t.start()
+    return t, box
+
+
+def test_light_predict_runs_while_a_fit_is_in_progress():
+    engine, backend = build()
+    engine.fit(*table(10, "x"), dataset_id="small")
+    backend.fit_gate = threading.Event()
+    fitting, _ = start(engine.fit, *table(100), dataset_id="big")
+    assert backend.fit_entered.wait(5)
+    # The fit is mid-flight; a small predict doesn't wait for it.
+    assert engine.predict("small", [[0.0]]).predictions == ["x"]
+    backend.fit_gate.set()
+    fitting.join(5)
+    assert engine._cache.holds("big")
+
+
+def test_large_predict_waits_for_the_fit():
+    engine, backend = build()
+    engine.fit(*table(10, "x"), dataset_id="small")
+    backend.fit_gate = threading.Event()
+    fitting, _ = start(engine.fit, *table(100), dataset_id="big")
+    assert backend.fit_entered.wait(5)
+    rows = [[0.0]] * (LIGHT_PREDICT_ROWS + 1)
+    predicting, box = start(engine.predict, "small", rows)
+    predicting.join(0.2)
+    assert predicting.is_alive() and backend.predict_calls == 0
+    backend.fit_gate.set()
+    fitting.join(5)
+    predicting.join(5)
+    assert len(box["result"].predictions) == LIGHT_PREDICT_ROWS + 1
+
+
+def test_restoring_a_spilled_context_waits_for_the_fit(tmp_path):
+    engine, backend = build(tmp_path)
+    engine.fit(*table(10, "x"), dataset_id="cold")
+    engine._cache.make_room(CAPACITY)  # spill it
+    assert not engine._cache.holds("cold") and engine._cache.has("cold")
+    backend.fit_gate = threading.Event()
+    fitting, _ = start(engine.fit, *table(100), dataset_id="big")
+    assert backend.fit_entered.wait(5)
+    predicting, box = start(engine.predict, "cold", [[0.0]])
+    predicting.join(0.2)
+    assert predicting.is_alive()  # a restore takes device memory
+    backend.fit_gate.set()
+    fitting.join(5)
+    predicting.join(5)
+    assert box["result"].predictions == ["x"]
+
+
+def test_evict_ahead_skips_a_context_in_use_then_takes_it_once_released(tmp_path):
+    engine, backend = build(tmp_path)
+    engine.fit(*table(600, "busy"), dataset_id="busy")  # 6,000 of 10,000 bytes
+    backend.predict_gate = threading.Event()
+    predicting, box = start(engine.predict, "busy", [[0.0]])
+    assert backend.predict_entered.wait(5)
+    # Needs 6,000 bytes: only evicting "busy" makes room, and it is pinned.
+    fitting, _ = start(engine.fit, *table(600), dataset_id="next")
+    fitting.join(0.2)
+    assert fitting.is_alive() and engine._cache.holds("busy")
+    backend.predict_gate.set()
+    predicting.join(5)
+    fitting.join(5)
+    assert box["result"].predictions == ["busy"]
+    assert engine._cache.holds("next") and not engine._cache.holds("busy")
+    assert engine._cache.has("busy")  # spilled, not lost
+
+
+def test_evict_ahead_prefers_unpinned_victims(tmp_path):
+    engine, backend = build(tmp_path)
+    engine.fit(*table(300, "busy"), dataset_id="busy")
+    engine.fit(*table(300, "idle"), dataset_id="idle")
+    backend.predict_gate = threading.Event()
+    predicting, _ = start(engine.predict, "busy", [[0.0]])
+    assert backend.predict_entered.wait(5)
+    # LRU would pick "busy" (older), but it's in use.
+    engine.fit(*table(500), dataset_id="new")
+    assert engine._cache.holds("busy") and not engine._cache.holds("idle")
+    backend.predict_gate.set()
+    predicting.join(5)
+
+
+def test_a_fit_that_runs_out_of_memory_retries_alone():
+    engine, backend = build()
+    backend.fail_next_fits = 1
+    engine.fit(*table(10, "x"), dataset_id="d")
+    assert backend.fit_calls == 1  # the failed attempt raised before counting
+    assert engine.predict("d", [[0.0]]).predictions == ["x"]
+
+
+def test_a_fit_failing_twice_raises():
+    engine, backend = build()
+    backend.fail_next_fits = 2
+    with pytest.raises(BackendComputeError):
+        engine.fit(*table(10), dataset_id="d")
+    assert not engine._cache.has("d")
+
+
+def test_stats_never_wait_for_a_fit():
+    engine, backend = build()
+    backend.fit_gate = threading.Event()
+    fitting, _ = start(engine.fit, *table(100), dataset_id="big")
+    assert backend.fit_entered.wait(5)
+    assert engine.stats().n_cached_contexts == 0  # the cache lock is free
+    backend.fit_gate.set()
+    fitting.join(5)

From da94ff7147d246abd4f91a19eeb975217f339be8 Mon Sep 17 00:00:00 2001
From: Amrit Krishnan <amrit110@gmail.com>
Date: Sat, 17 Oct 2026 20:15:56 -0400
Subject: [PATCH 18/39] fix: review findings on the worker-process and
 scheduling commits

1. Deadlock: a large predict took _fit_lock then the host-memory slot
   while fits take the slot first. One order everywhere now: host slot,
   _fit_lock, _predict_lock (cache and release locks are leaves).
2. A hot context kept pinned by queued predicts made the fit go ahead
   without room (cache over capacity). When only in-use contexts stand in
   the way, the fit holds predicts off and evicts them; predicts hold no
   context reference while waiting and re-check residency once they have
   the device (restoring if needed). put() warns if it ever exceeds
   capacity.
3. ProcessShard: errors are checked with a pickle round trip; an
   unreadable reply retires the worker instead of killing the reader
   thread silently; worker-side send failures are reported; quick calls
   (estimate, report) have timeouts; argument pickling errors don't leak
   futures; a dead/unresponsive worker is skipped by placement.
4. An evict during a fit freed memory inside the fit's allocator-delta
   measurement: removed contexts are parked until the fit has measured.
5. Fit OOM retry (and predict halving) now run outside the except block,
   after gc + empty_cache, with room made again.
6. A short runtime fit no longer sets a width-dependent model's rate when
   long fits that wide exist (runtime long fits still supersede).
7. A device OOM while restoring keeps the spilled copy (was deleted as
   corrupt); restores make room first and retry with predicts held off.
8. /readyz maps worker death/timeouts to 503.
9. put() keeps a replaced id visible to has() while it makes room.
10. Spill reads/writes run outside the cache lock (light predicts no
    longer wait on multi-GB I/O); restores wait for in-transit ids.
12. sweep_engine/test use built.default/engine.backend_for; the pressure
    bench reads its queue before join.

Regression tests for 1, 2, 3, 4, 6, 7 and 9 (the deadlock one fails on
the old code by hanging; 4, 7 and 9 fail on it outright).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 benchmarks/bench_memory_pressure.py      |  15 +-
 benchmarks/sweep_engine.py               |   2 +-
 src/tabctx/backends/_cuda.py             |  11 +
 src/tabctx/cache/manager.py              | 220 ++++++++++--------
 src/tabctx/cache/spill.py                |  37 +++-
 src/tabctx/engine.py                     | 271 +++++++++++++++++------
 src/tabctx/memory/requirement.py         |  56 ++++-
 src/tabctx/serve/app.py                  |   6 +-
 src/tabctx/serve/workers.py              |  72 +++++-
 tests/integration/test_tabpfn_backend.py |   2 +-
 tests/unit/test_engine_scheduling.py     |  97 ++++++++
 tests/unit/test_requirement.py           |  19 ++
 tests/unit/test_spill.py                 |  19 ++
 tests/unit/test_workers.py               |  12 +
 14 files changed, 641 insertions(+), 198 deletions(-)

diff --git a/benchmarks/bench_memory_pressure.py b/benchmarks/bench_memory_pressure.py
index 212608e..2ef76db 100644
--- a/benchmarks/bench_memory_pressure.py
+++ b/benchmarks/bench_memory_pressure.py
@@ -141,14 +141,19 @@ def main() -> int:
         queue = ctx.Queue()
         child = ctx.Process(target=_measure_shape, args=(args, rows, features, queue))
         child.start()
+        # Read before join: a child blocks exiting until its queued output
+        # is consumed.
+        records = None
+        while records is None and (child.is_alive() or not queue.empty()):
+            try:
+                records = queue.get(timeout=5)
+            except Exception:  # noqa: BLE001 -- queue.Empty; keep waiting
+                continue
         child.join()
-        records = (
-            queue.get()
-            if not queue.empty()
-            else [
+        if records is None:
+            records = [
                 {"rows": rows, "features": features, "status": f"exit {child.exitcode}"}
             ]
-        )
         for record in records:
             line = json.dumps(record)
             print(line, flush=True)
diff --git a/benchmarks/sweep_engine.py b/benchmarks/sweep_engine.py
index 354527c..babcfde 100644
--- a/benchmarks/sweep_engine.py
+++ b/benchmarks/sweep_engine.py
@@ -57,7 +57,7 @@ def main() -> int:
                 rec["status"] = "OOM"
                 rec["detail"] = str(e)[:160]
             rec["fit_s"] = round(time.perf_counter() - t0, 2)
-            backend = built.backend
+            backend = built.engine.backend_for(built.default)
             if rec["status"] == "ok":
                 rec["peak_gib"] = round((backend.fit_peak_bytes_hint() or 0) / 2**30, 2)
                 rec["resident_gib"] = round(
diff --git a/src/tabctx/backends/_cuda.py b/src/tabctx/backends/_cuda.py
index b414acb..67a4a12 100644
--- a/src/tabctx/backends/_cuda.py
+++ b/src/tabctx/backends/_cuda.py
@@ -29,6 +29,17 @@ def ensure_initialized(device: str) -> None:
         torch.cuda.init()
 
 
+def release_cached_device_memory() -> None:
+    """Return the CUDA caching allocator's unused blocks (a retry after an
+    OOM then sees an unfragmented device). A no-op without torch or CUDA."""
+    try:
+        import torch
+    except ImportError:
+        return
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+
 def device_scope(device: str) -> contextlib.AbstractContextManager[Any]:
     """Make ``device`` torch's current CUDA device for the block; a no-op
     on CPU."""
diff --git a/src/tabctx/cache/manager.py b/src/tabctx/cache/manager.py
index 0ec56f8..176d0c6 100644
--- a/src/tabctx/cache/manager.py
+++ b/src/tabctx/cache/manager.py
@@ -20,6 +20,7 @@
 from __future__ import annotations
 
 import contextlib
+import logging
 import threading
 import time
 from collections.abc import Iterator
@@ -30,6 +31,8 @@
 from tabctx.errors import CacheCapacityError
 from tabctx.types import EngineStats, Task
 
+logger = logging.getLogger(__name__)
+
 
 @dataclass
 class CachedContext:
@@ -69,48 +72,82 @@ def __init__(
         self._used_bytes = 0
         self._spill = spill_store
         self._lock = threading.RLock()
-        # Ids moving between tiers (being spilled, or restored), which
-        # neither tier lists until the I/O finishes. Lock-free readers
-        # (has) must still see them: a multi-device router that saw
-        # neither would answer 404 for a context that is merely moving.
-        self._in_transit: set[str] = set()
+        # Ids moving between tiers (being spilled, restored, or replaced by a
+        # re-fit), which neither tier lists until the move finishes; counted,
+        # since moves can nest (a restore's put). Lock-free readers (has)
+        # must still see them: a multi-device router that saw neither would
+        # answer 404 for a context that is merely moving. Spill I/O runs
+        # outside the lock, so bookkeeping (and light predicts) never wait
+        # on a multi-GB disk write.
+        self._in_transit: dict[str, int] = {}
+        self._settled = threading.Condition(self._lock)
         # Contexts a call is using right now (see pinned()): never chosen as
-        # eviction victims, so a fit's evict-ahead can't spill a context out
-        # from under a predict running beside it -- the predict would keep
-        # the GPU memory alive while the books count it free.
+        # policy eviction victims, so a fit's evict-ahead can't spill a
+        # context out from under a predict running beside it -- the predict
+        # would keep the GPU memory alive while the books count it free.
         self._pins: dict[str, int] = {}
-        self._unpinned = threading.Condition(self._lock)
+
+    # ---- transit bookkeeping (caller holds the lock) -----------------
+
+    def _transit_add(self, dataset_id: str) -> None:
+        self._in_transit[dataset_id] = self._in_transit.get(dataset_id, 0) + 1
+
+    def _transit_discard(self, dataset_id: str) -> None:
+        left = self._in_transit.get(dataset_id, 0) - 1
+        if left > 0:
+            self._in_transit[dataset_id] = left
+        else:
+            self._in_transit.pop(dataset_id, None)
+            self._settled.notify_all()
+
+    def wait_settled(self, dataset_id: str, timeout_s: float = 600.0) -> bool:
+        """Wait until an id is no longer moving between tiers (a spill
+        write, restore or replacement in progress). False on timeout."""
+        with self._settled:
+            return self._settled.wait_for(
+                lambda: dataset_id not in self._in_transit, timeout_s
+            )
 
     def get(self, dataset_id: str) -> CachedContext | None:
+        """The resident context, else a spilled one reloaded and re-admitted
+        (its displacements may spill in turn). The reload runs outside the
+        lock; callers serialize restores of one device (the engine does so
+        under its fit lock)."""
         with self._lock:
             entry = self._entries.get(dataset_id)
             if entry is not None or self._spill is None:
                 return entry
-            # Miss with a spill tier: try reloading. Re-admitting evicts
-            # via the normal policy to make room (RLock, so the nested
-            # put() is fine); what it displaces may itself spill.
-            self._in_transit.add(dataset_id)
+            self._transit_add(dataset_id)
+        try:
+            restored = self._spill.load(dataset_id)
+            if restored is None:
+                return None
+            restored.last_accessed_at = time.monotonic()
             try:
-                restored = self._spill.load(dataset_id)
-                if restored is None:
-                    return None
-                restored.last_accessed_at = time.monotonic()
-                try:
-                    self.put(restored)
-                except CacheCapacityError:
-                    # Spilled under a larger budget than this process has
-                    # (e.g. TABCTX_GPU_MEMORY_FRACTION lowered across a
-                    # restart). Nothing can re-admit it; treat as a miss so
-                    # the caller re-fits instead of the request 500ing.
-                    return None
-                return restored
-            finally:
-                self._in_transit.discard(dataset_id)
+                self.put(restored)
+            except CacheCapacityError:
+                # Spilled under a larger budget than this process has (e.g.
+                # TABCTX_GPU_MEMORY_FRACTION lowered across a restart).
+                # Nothing can re-admit it; treat as a miss so the caller
+                # re-fits instead of the request 500ing.
+                return None
+            return restored
+        finally:
+            with self._lock:
+                self._transit_discard(dataset_id)
+
+    def spilled_size(self, dataset_id: str) -> int | None:
+        """A spilled context's est_bytes, without loading it (None if not
+        spilled): what restoring it will occupy."""
+        if self._spill is None:
+            return None
+        return self._spill.est_bytes(dataset_id)
 
     @contextlib.contextmanager
     def pinned(self, dataset_id: str) -> Iterator[None]:
-        """Keep a context out of eviction's reach for the duration (explicit
-        evict() and re-fits still replace it: caller intent)."""
+        """Keep a context out of policy eviction's reach for the duration
+        (explicit evict(), re-fits and evict_one(ignore_pins=True) still
+        take it: caller intent)."""
         with self._lock:
             self._pins[dataset_id] = self._pins.get(dataset_id, 0) + 1
         try:
@@ -122,18 +159,9 @@ def pinned(self, dataset_id: str) -> Iterator[None]:
                     self._pins[dataset_id] = left
                 else:
                     del self._pins[dataset_id]
-                    self._unpinned.notify_all()
-
-    def wait_for_unpin(self, timeout_s: float) -> bool:
-        """Wait until some pinned context is released. False if nothing is
-        pinned (nothing to wait for) or on timeout."""
-        with self._unpinned:
-            if not self._pins:
-                return False
-            before = sum(self._pins.values())
-            return self._unpinned.wait_for(
-                lambda: sum(self._pins.values()) < before, timeout_s
-            )
+
+    def has_pins(self) -> bool:
+        return bool(self._pins)
 
     def peek(self, dataset_id: str) -> CachedContext | None:
         """The resident context, or None -- never reloads a spilled one (a
@@ -155,54 +183,53 @@ def _remove(self, dataset_id: str) -> CachedContext | None:
             self._used_bytes -= entry.est_bytes
         return entry
 
-    def _evict_victim(self) -> str | None:
-        """Policy-chosen eviction, spilling when a tier is attached.
-        Caller holds the lock."""
-        victim_id = self._policy.select_victim(
-            [e for e in self._entries.values() if e.dataset_id not in self._pins]
-        )
-        if victim_id is None:
-            return None
-        if self._spill is None:
-            self._remove(victim_id)
-            return victim_id
-        self._in_transit.add(victim_id)
-        try:
+    def evict_one(self, ignore_pins: bool = False) -> str | None:
+        """Evict a single victim chosen by the policy (spilling it when a
+        spill tier is attached). Returns the evicted dataset_id, or None if
+        nothing is evictable. Used by the engine's evict-ahead-of-fit path
+        to convert resident bytes into transient fit headroom; ignore_pins
+        is for callers that have stopped every predict first."""
+        with self._lock:
+            candidates = [
+                e
+                for e in self._entries.values()
+                if ignore_pins or e.dataset_id not in self._pins
+            ]
+            victim_id = self._policy.select_victim(candidates)
+            if victim_id is None:
+                return None
             victim = self._remove(victim_id)
-            self._spill.spill(victim)  # best-effort by contract
+            if self._spill is None:
+                return victim_id
+            self._transit_add(victim_id)
+        try:
+            self._spill.spill(victim)  # best-effort by contract; no lock held
         finally:
-            self._in_transit.discard(victim_id)
+            with self._lock:
+                self._transit_discard(victim_id)
         return victim_id
 
     def make_room(self, needed_bytes: int) -> list[str]:
         """Evict entries (via the configured policy) until at least
         needed_bytes is free, or there's nothing left to evict. Does not
         raise -- callers must check free_bytes afterward if they need a
-        guarantee (put() does this for you)."""
+        guarantee."""
         evicted: list[str] = []
-        with self._lock:
-            while self.free_bytes < needed_bytes:
-                victim_id = self._evict_victim()
-                if victim_id is None:
-                    break
-                evicted.append(victim_id)
+        while self.free_bytes < needed_bytes:
+            victim_id = self.evict_one()
+            if victim_id is None:
+                break
+            evicted.append(victim_id)
         return evicted
 
-    def evict_one(self) -> str | None:
-        """Evict a single victim chosen by the policy (spilling it when a
-        spill tier is attached). Returns the evicted dataset_id, or None
-        if the cache is empty. Used by the engine's evict-ahead-of-fit
-        path to convert resident bytes into transient fit headroom."""
-        with self._lock:
-            return self._evict_victim()
-
     def put(self, context: CachedContext) -> list[str]:
         """Insert a context, evicting via the policy to make room first.
         Returns the list of evicted dataset_ids. Raises CacheCapacityError
         if context.est_bytes alone exceeds total capacity -- no amount of
         eviction would ever make room for it. Overwriting an existing
         dataset_id replaces it in both tiers (a stale spilled copy must
-        never outlive the context that superseded it)."""
+        never outlive the context that superseded it); the id stays
+        visible to has() throughout."""
         if context.est_bytes > self._capacity_bytes:
             raise CacheCapacityError(
                 f"context for {context.dataset_id!r} needs "
@@ -210,30 +237,47 @@ def put(self, context: CachedContext) -> list[str]:
                 f"{self._capacity_bytes} bytes"
             )
         with self._lock:
+            self._transit_add(context.dataset_id)
             self._remove(context.dataset_id)
+        try:
             if self._spill is not None:
                 self._spill.delete(context.dataset_id)
             evicted = self.make_room(context.est_bytes)
-            self._entries[context.dataset_id] = context
-            self._used_bytes += context.est_bytes
+            with self._lock:
+                self._entries[context.dataset_id] = context
+                self._used_bytes += context.est_bytes
+                if self._used_bytes > self._capacity_bytes:
+                    # Only pinned contexts were left to evict: the books
+                    # stay truthful (the memory is in use), but say so.
+                    logger.warning(
+                        "cache over capacity after inserting %r: %d of %d bytes "
+                        "(contexts in use could not be evicted)",
+                        context.dataset_id,
+                        self._used_bytes,
+                        self._capacity_bytes,
+                    )
             return evicted
+        finally:
+            with self._lock:
+                self._transit_discard(context.dataset_id)
 
-    def evict(self, dataset_id: str) -> None:
-        """Drop a context from every tier. Never spills: this is caller
-        intent (fit_predict's one-shot cleanup, an explicit delete), and
-        a spilled copy left behind would silently resurrect on the next
-        get()."""
+    def evict(self, dataset_id: str) -> CachedContext | None:
+        """Drop a context from every tier and return the resident one (the
+        caller decides when its memory may be released). Never spills: this
+        is caller intent (fit_predict's one-shot cleanup, an explicit
+        delete), and a spilled copy left behind would silently resurrect on
+        the next get()."""
         with self._lock:
-            self._remove(dataset_id)
-            if self._spill is not None:
-                self._spill.delete(dataset_id)
+            removed = self._remove(dataset_id)
+        if self._spill is not None:
+            self._spill.delete(dataset_id)
+        return removed
 
     # ---- lock-free snapshots ---------------------------------------
-    # The engine holds this cache's lock across whole GPU calls, so a
-    # reader that only needs bookkeeping (a multi-device router choosing
-    # a device, /readyz) must not take it: it would wait for that GPU's
-    # work to finish. Single int / dict-membership reads are atomic under
-    # the GIL; these are snapshots, and every mutation stays locked.
+    # Readers that only need bookkeeping (a multi-device router choosing a
+    # device, /readyz) never take the lock. Single int / dict-membership
+    # reads are atomic under the GIL; these are snapshots, and every
+    # mutation stays locked.
 
     def holds(self, dataset_id: str) -> bool:
         """Resident in the primary tier (snapshot, lock-free)."""
diff --git a/src/tabctx/cache/spill.py b/src/tabctx/cache/spill.py
index 71ab405..2307cd0 100644
--- a/src/tabctx/cache/spill.py
+++ b/src/tabctx/cache/spill.py
@@ -50,6 +50,16 @@
 DEFAULT_SPILL_CAPACITY_BYTES = 50 * 1024**3
 
 
+def _is_out_of_memory(error: BaseException) -> bool:
+    """A device OOM (torch's, or a backend's wrapper of it), without
+    importing torch."""
+    from tabctx.errors import BackendComputeError
+
+    return (
+        isinstance(error, BackendComputeError) or "OutOfMemory" in type(error).__name__
+    )
+
+
 class DiskSpillStore:
     def __init__(
         self,
@@ -127,7 +137,9 @@ def spill(self, context: CachedContext) -> bool:
     def load(self, dataset_id: str) -> CachedContext | None:
         """Reload a spilled context (removing it from the spill tier --
         it's going back into the primary cache). None if absent or
-        unreadable."""
+        unreadable. A device out-of-memory while deserializing onto the GPU
+        is not corruption: the files are kept and the error propagates, so
+        a retry (or a later request) can still restore it."""
         meta_path, payload_path = self._paths(dataset_id)
         try:
             meta = json.loads(meta_path.read_text())
@@ -135,13 +147,14 @@ def load(self, dataset_id: str) -> CachedContext | None:
                 meta["backend_name"], (pickle.dumps, pickle.loads)
             )
             payload = loads(payload_path.read_bytes())
-        except Exception:  # noqa: BLE001 -- absent or corrupt == miss
-            return None
-        finally:
-            # Success: the context moves back to the primary cache.
-            # Failure: drop the unreadable files rather than retrying
-            # them forever.
+        except Exception as e:  # noqa: BLE001 -- absent or corrupt == miss
+            if _is_out_of_memory(e):
+                raise
+            # Unreadable: drop the files rather than retrying them forever.
             self.delete(dataset_id)
+            return None
+        # The context moves back to the primary cache.
+        self.delete(dataset_id)
         return CachedContext(
             dataset_id=meta["dataset_id"],
             backend_name=meta["backend_name"],
@@ -154,6 +167,16 @@ def load(self, dataset_id: str) -> CachedContext | None:
             feature_names=meta.get("feature_names"),
         )
 
+    def est_bytes(self, dataset_id: str) -> int | None:
+        """A spilled context's recorded size, from its metadata only."""
+        if dataset_id not in self._index:
+            return None
+        meta_path, _ = self._paths(dataset_id)
+        try:
+            return int(json.loads(meta_path.read_text())["est_bytes"])
+        except (OSError, ValueError, KeyError):
+            return None
+
     def contains(self, dataset_id: str) -> bool:
         """Whether a spilled copy exists. Lock-free snapshot (a dict
         membership read is atomic), for routing decisions that must not
diff --git a/src/tabctx/engine.py b/src/tabctx/engine.py
index d98f902..5f66bb4 100644
--- a/src/tabctx/engine.py
+++ b/src/tabctx/engine.py
@@ -7,9 +7,12 @@
 from __future__ import annotations
 
 import contextlib
+import gc
 import threading
 import uuid
+from collections.abc import Iterator
 
+from tabctx.backends._cuda import release_cached_device_memory
 from tabctx.backends.base import TabularICLBackend
 from tabctx.cache.manager import CachedContext, ContextCacheManager
 from tabctx.chunking import choose_chunk_rows, split_rows
@@ -30,9 +33,17 @@
 # predicts, and predicts that must first restore a spilled context, need
 # device memory a running fit may be using, so they wait their turn.
 LIGHT_PREDICT_ROWS = 2048
-# How long a fit's evict-ahead waits for pinned (in-use) contexts to be
-# released before trying its luck with the headroom it has.
-UNPIN_WAIT_S = 30.0
+# A predict whose context is evicted again (under memory pressure) between
+# its restore and getting the device tries this many times.
+_PREDICT_ATTEMPTS = 3
+
+
+def _not_found(dataset_id: str) -> DatasetNotFoundError:
+    return DatasetNotFoundError(
+        f"no cached context for dataset_id={dataset_id!r} -- it was "
+        "never fit, was evicted, or was lost to a replica restart "
+        "(the cache has no durability across process restarts)"
+    )
 
 
 def _rect_shape(X: ArrayLike) -> tuple[int, int] | None:
@@ -106,16 +117,19 @@ def _predict_splitting_on_oom(
     """Predict, halving the chunk on a CUDA OOM until it fits (or one row
     still doesn't). Chunk sizing estimates can't see everything on a nearly
     full device -- e.g. a compressed kv cache's per-layer decompression --
-    and a model may pick a single batch with no OOM retry of its own."""
+    and a model may pick a single batch with no OOM retry of its own. The
+    halves run after the except block, whose traceback would otherwise keep
+    the failed attempt's tensors alive."""
     try:
         return [backend.predict(payload, X_test, return_proba=return_proba)]
     except BackendComputeError:
         if len(X_test) <= 1:
             raise
-        half = len(X_test) // 2
-        return _predict_splitting_on_oom(
-            backend, payload, X_test[:half], return_proba
-        ) + _predict_splitting_on_oom(backend, payload, X_test[half:], return_proba)
+    gc.collect()
+    half = len(X_test) // 2
+    return _predict_splitting_on_oom(
+        backend, payload, X_test[:half], return_proba
+    ) + _predict_splitting_on_oom(backend, payload, X_test[half:], return_proba)
 
 
 class TabctxEngine:
@@ -166,11 +180,17 @@ def __init__(
         # large predicts. _predict_lock runs one GPU predict at a time, so
         # a light predict (see LIGHT_PREDICT_ROWS) overlaps a running fit
         # rather than queueing behind it. The cache's own lock only guards
-        # bookkeeping. Lock order: _fit_lock, then _predict_lock.
+        # bookkeeping. Lock order: host slot, _fit_lock, _predict_lock (the
+        # cache lock and _release_lock are leaves).
         self._fit_lock = threading.Lock()
         # Re-entrant: a fit retrying alone holds it while the backend's
         # measurement takes it again (see _cuda.QuiescedMeasurement).
         self._predict_lock = threading.RLock()
+        # Contexts removed while a fit measures itself, freed after (see
+        # _measuring); the release lock orders frees against measurement.
+        self._release_lock = threading.Lock()
+        self._fitting = False
+        self._parked: list[CachedContext] = []
         for b in backends.values():
             lend = getattr(b, "set_measurement_quiesce", None)
             if callable(lend):
@@ -194,6 +214,10 @@ def backend_names(self) -> list[str]:
     def default_backend(self) -> str:
         return self._default_backend
 
+    def backend_for(self, backend: str | None = None) -> TabularICLBackend:
+        """The backend instance serving a model (the default when None)."""
+        return self._resolve_backend(backend)[1]
+
     def estimator_for(self, backend: str | None = None) -> MemoryEstimator:
         return self._estimators[backend or self._default_backend]
 
@@ -288,17 +312,19 @@ def fit(
         estimator = self._estimators[backend_name]
 
         resolved_id = dataset_id or str(uuid.uuid4())
-        # Host memory first: a large fit waiting for a host slot must not
-        # hold this device meanwhile.
+        # Lock order everywhere: host slot, _fit_lock, _predict_lock. Host
+        # memory first, so a large fit waiting for a host slot doesn't hold
+        # this device meanwhile.
         with self._host_slot(chosen, n_train, n_features), self._fit_lock:
             estimated = estimator.estimate_bytes(n_train, 0, n_features)
             self._admit(estimator, estimated, n_train, n_features)
-            payload = self._fit_backend(chosen, X, y, task)
-            est_bytes = chosen.context_bytes_hint(n_train, n_features)
-            if est_bytes is None:
-                est_bytes = estimated
-            else:
-                self._learn(chosen, estimator, n_train, n_features, est_bytes)
+            with self._measuring():
+                payload = self._fit_backend(chosen, X, y, task, estimator, estimated)
+                est_bytes = chosen.context_bytes_hint(n_train, n_features)
+                if est_bytes is None:
+                    est_bytes = estimated
+                else:
+                    self._learn(chosen, estimator, n_train, n_features, est_bytes)
             context = CachedContext(
                 dataset_id=resolved_id,
                 backend_name=backend_name,
@@ -329,8 +355,7 @@ def _admit(
         the need fits -- without this, a filling replica would reject fits
         that are perfectly safe once a cold context steps aside (proven on
         real hardware: a fourth ~13GB-peak fit OOMed an A100 with three
-        ~8GB contexts resident). Contexts a predict is using are pinned and
-        skipped; if only pinned ones stand in the way, wait for them."""
+        ~8GB contexts resident)."""
         # Feasibility FIRST, against an EMPTY cache's headroom: a fit that
         # can never be admitted must be rejected before evicting anything --
         # otherwise one oversized request drains every tenant's context on
@@ -345,28 +370,73 @@ def _admit(
                 f"{max_headroom} byte maximum "
                 "admission headroom (even with an empty cache)"
             )
-        with self._cache.lock:
-            while True:
-                headroom = self._headroom(estimator, self._cache.stats().used_bytes)
-                if estimated <= headroom:
-                    return
-                if self._cache.evict_one() is not None:
-                    continue
-                # Only pinned contexts left (or none at all): wait for a
-                # predict to finish, else go ahead with what there is.
-                if not self._cache.wait_for_unpin(UNPIN_WAIT_S):
-                    return
+        self._make_room(estimator, estimated)
+
+    def _make_room(self, estimator: MemoryEstimator, needed: int) -> None:
+        """Evict until `needed` bytes of headroom are free, or nothing is
+        left to evict. Caller holds _fit_lock (and no cache or predict
+        lock). Contexts a predict is using are pinned and evicted last:
+        when only they stand in the way, predicts are held off (no predict
+        is then running, and queued ones re-check residency once they get
+        the device) and those contexts go too."""
+
+        def short() -> bool:
+            return needed > self._headroom(estimator, self._cache.stats().used_bytes)
+
+        while short():
+            if self._cache.evict_one() is None:
+                break
+        if short() and self._cache.has_pins():
+            with self._predict_lock:
+                while short() and self._cache.evict_one(ignore_pins=True) is not None:
+                    pass
+
+    @contextlib.contextmanager
+    def _measuring(self) -> Iterator[None]:
+        """While a fit measures its footprint as an allocator delta, no
+        other context's memory may be freed on this device: evicted
+        payloads are parked (see _release) and dropped afterwards."""
+        with self._release_lock:
+            self._fitting = True
+        try:
+            yield
+        finally:
+            with self._release_lock:
+                self._fitting = False
+                parked, self._parked = self._parked, []
+            del parked
+
+    def _release(self, context: CachedContext | None) -> None:
+        """Drop the caller's last reference to a removed context now, or
+        park it until the running fit has measured itself. The caller must
+        not keep another reference."""
+        with self._release_lock:
+            if context is not None and self._fitting:
+                self._parked.append(context)
+            context = None
 
     def _fit_backend(
-        self, backend: TabularICLBackend, X: ArrayLike, y: ArrayLike, task: Task
+        self,
+        backend: TabularICLBackend,
+        X: ArrayLike,
+        y: ArrayLike,
+        task: Task,
+        estimator: MemoryEstimator,
+        estimated: int,
     ) -> object:
-        """Fit, retrying once with predicts held off if it ran out of
-        device memory while sharing the device with one."""
+        """Fit, retrying once if it ran out of device memory while sharing
+        the device with a predict: after the failed attempt's tensors are
+        freed (outside the except block, whose traceback would keep them
+        alive), with predicts held off and room made again."""
         try:
             return backend.fit(X, y, task)
         except BackendComputeError:
-            with self._predict_lock:
-                return backend.fit(X, y, task)
+            pass
+        gc.collect()
+        release_cached_device_memory()
+        with self._predict_lock:
+            self._make_room(estimator, estimated)
+            return backend.fit(X, y, task)
 
     def _learn(
         self,
@@ -412,14 +482,47 @@ def _host_slot(self, backend: TabularICLBackend, n_rows: int, n_features: int):
         )
         return self._host_guard.reserve(large)
 
+    def _restore(self, dataset_id: str) -> None:
+        """Bring a spilled context back onto the device (DatasetNotFoundError
+        if it is neither resident nor spilled). Restores grow residency, so
+        they take _fit_lock and make room for the context first; one that
+        runs out of memory beside a predict retries with predicts held off
+        (the spilled copy survives a device OOM)."""
+        with self._fit_lock:
+            # A context mid-spill is in neither tier until the write lands.
+            self._cache.wait_settled(dataset_id)
+            if self._cache.peek(dataset_id) is not None:
+                return
+            size = self._cache.spilled_size(dataset_id)
+            if size is None:
+                raise _not_found(dataset_id)
+            estimator = self._estimators[self._default_backend]
+            self._make_room(estimator, size)
+            try:
+                restored = self._cache.get(dataset_id)
+            except BackendComputeError:
+                restored = None
+                retry = True
+            else:
+                retry = False
+            if retry:
+                gc.collect()
+                release_cached_device_memory()
+                with self._predict_lock:
+                    restored = self._cache.get(dataset_id)
+            if restored is None:
+                raise _not_found(dataset_id)
+
     def feature_names(self, dataset_id: str) -> list[str] | None:
         """Column names recorded at fit() for a cached context, or None
         when the context is unknown or was fit without names. Reloads a
         spilled context as a side effect, exactly as predict() would."""
+        if self._cache.peek(dataset_id) is None:
+            try:
+                self._restore(dataset_id)
+            except DatasetNotFoundError:
+                return None
         context = self._cache.peek(dataset_id)
-        if context is None:
-            with self._fit_lock:  # a restore takes device memory
-                context = self._cache.get(dataset_id)
         return context.feature_names if context is not None else None
 
     def predict(
@@ -432,69 +535,89 @@ def predict(
         which the naive one-shot wrapper this library replaces did not do.
 
         A light predict (a resident context, at most LIGHT_PREDICT_ROWS test
-        rows) runs beside a fit on this device; anything else waits for the
-        device's fits.
+        rows) runs beside a fit on this device; larger ones wait for the
+        device's fits, and a spilled context is restored first.
         """
+        exclusive = len(X_test) > LIGHT_PREDICT_ROWS
         with self._cache.pinned(dataset_id):
-            context = self._cache.peek(dataset_id)
-            if context is not None and len(X_test) <= LIGHT_PREDICT_ROWS:
-                return self._predict_on(context, X_test, return_proba)
-        with self._fit_lock, self._cache.pinned(dataset_id):
-            context = self._cache.get(dataset_id)  # restores a spilled one
-            if context is None:
-                raise DatasetNotFoundError(
-                    f"no cached context for dataset_id={dataset_id!r} -- it was "
-                    "never fit, was evicted, or was lost to a replica restart "
-                    "(the cache has no durability across process restarts)"
+            for _ in range(_PREDICT_ATTEMPTS):
+                if self._cache.peek(dataset_id) is None:
+                    self._restore(dataset_id)
+                outcome = self._predict_resident(
+                    dataset_id, X_test, return_proba, exclusive
                 )
-            return self._predict_on(context, X_test, return_proba)
+                if outcome is not None:
+                    return outcome
+        raise BackendComputeError(
+            f"context {dataset_id!r} was evicted again before its predict "
+            f"could run ({_PREDICT_ATTEMPTS} attempts); the device is under "
+            "memory pressure -- retry"
+        )
 
-    def _predict_on(
-        self, context: CachedContext, X_test: ArrayLike, return_proba: bool
-    ) -> PredictOutcome:
-        """Run a predict against a context the caller has pinned."""
+    def _predict_resident(
+        self, dataset_id: str, X_test: ArrayLike, return_proba: bool, exclusive: bool
+    ) -> PredictOutcome | None:
+        """Run a predict if the context is resident once this call has the
+        device (None if it was evicted meanwhile). Holds no reference to the
+        context while waiting for locks, so an eviction can really free it."""
+        context = self._cache.peek(dataset_id)
+        if context is None:
+            return None
         _validate_predict_input(X_test, context.n_features)
         # Dispatch to whichever backend fit this context (multi-backend
         # deployments serve several models over one cache).
         backend = self._backends.get(context.backend_name)
         if backend is None:
             raise DatasetNotFoundError(
-                f"context for {context.dataset_id!r} was fit by backend "
+                f"context for {dataset_id!r} was fit by backend "
                 f"{context.backend_name!r}, which this deployment no longer "
                 "serves -- re-fit with one of "
                 f"{sorted(self._backends)}"
             )
         estimator = self._estimators[context.backend_name]
-
-        n_test = len(X_test)
+        n_features = context.n_features
         # Chunk against the usage-aware transient headroom (what the
         # device can actually take on top of resident contexts), not the
         # cache-capacity ceiling -- the same quantity fit admission uses.
         chunk_rows = choose_chunk_rows(
             estimator,
             context.n_train,
-            context.n_features,
+            n_features,
             self._headroom(estimator, self._cache.stats_nowait().used_bytes),
         )
+        context = None
+        n_test = len(X_test)
         chunks = split_rows(X_test, chunk_rows) if chunk_rows < n_test else [X_test]
+        largest_chunk = max(len(c) for c in chunks)
 
         all_predictions: list = []
         all_probabilities: list[list[float]] = []
         classes: list[str] | None = None
-        largest_chunk = max(len(c) for c in chunks)
         with (
-            self._host_slot(backend, largest_chunk, context.n_features),
+            self._host_slot(backend, largest_chunk, n_features),
+            self._fit_lock if exclusive else contextlib.nullcontext(),
             self._predict_lock,
         ):
-            self._cache.touch(context.dataset_id)
-            for chunk in chunks:
-                for outcome in _predict_splitting_on_oom(
-                    backend, context.payload, chunk, return_proba
-                ):
-                    all_predictions.extend(outcome.predictions)
-                    if return_proba:
-                        all_probabilities.extend(outcome.probabilities or [])
-                        classes = outcome.classes or classes
+            context = self._cache.peek(dataset_id)
+            if context is None:
+                return None
+            self._cache.touch(dataset_id)
+            try:
+                for chunk in chunks:
+                    for outcome in _predict_splitting_on_oom(
+                        backend, context.payload, chunk, return_proba
+                    ):
+                        all_predictions.extend(outcome.predictions)
+                        if return_proba:
+                            all_probabilities.extend(outcome.probabilities or [])
+                            classes = outcome.classes or classes
+            finally:
+                # Still under the predict lock: if the context was evicted
+                # while this ran, its memory is freed now -- atomic with
+                # respect to a fit's measurements -- or parked.
+                if self._cache.peek(dataset_id) is not context:
+                    self._release(context)
+                context = None
 
         return PredictOutcome(
             predictions=all_predictions,
@@ -521,7 +644,13 @@ def fit_predict(
             self.evict(dataset_id)
 
     def evict(self, dataset_id: str) -> None:
-        self._cache.evict(dataset_id)
+        # Under the release lock, so the memory is freed either before a
+        # fit starts measuring or after it is done (see _measuring).
+        with self._release_lock:
+            removed = self._cache.evict(dataset_id)
+            if removed is not None and self._fitting:
+                self._parked.append(removed)
+            removed = None
 
     def stats(self) -> EngineStats:
         return self._cache.stats()
diff --git a/src/tabctx/memory/requirement.py b/src/tabctx/memory/requirement.py
index 6dad7a2..449d157 100644
--- a/src/tabctx/memory/requirement.py
+++ b/src/tabctx/memory/requirement.py
@@ -151,15 +151,9 @@ def record(self, n_train: int, n_features: int, resident_bytes: int) -> None:
     def _slope(
         self, points: list[tuple[int, int, int]], n_features: int
     ) -> float | None:
-        rates = [(n, d, res / n) for n, d, res in points]
-        long_fits = [(d, r) for n, d, r in rates if n >= self._min_rows]
-        if self._width_dependent:
-            wide = [d for _, d, _ in rates if d >= n_features]
-            if not wide:
-                return None
-            at = min(wide)
-            long_at = [r for d, r in long_fits if d == at]
-            return max(long_at or [r for _, d, r in rates if d == at])
+        """Width-independent rate: steepest long fit at the narrowest
+        measured width >= n_features, else among all long fits."""
+        long_fits = [(d, res / n) for n, d, res in points if n >= self._min_rows]
         if not long_fits:
             return None
         widths = [d for d, _ in long_fits if d >= n_features]
@@ -167,12 +161,50 @@ def _slope(
             return max(rate for _, rate in long_fits)
         return max(rate for d, rate in long_fits if d == min(widths))
 
+    def _width_dependent_slope(
+        self,
+        runtime: list[tuple[int, int, int]],
+        n_features: int,
+    ) -> float | None:
+        """Rate from fits at least as wide as the query only: long fits at the
+        narrowest width >= n_features that has any (this replica's own
+        superseding the calibration's), else -- only when no long fit that
+        wide exists anywhere -- short fits at the narrowest such width,
+        whose fixed overhead overstates the rate (conservative)."""
+
+        def long_rate(points: list[tuple[int, int, int]]) -> float | None:
+            wide = [
+                (d, res / n)
+                for n, d, res in points
+                if n >= self._min_rows and d >= n_features
+            ]
+            if not wide:
+                return None
+            at = min(d for d, _ in wide)
+            return max(r for d, r in wide if d == at)
+
+        rate = long_rate(runtime)
+        if rate is None:
+            rate = long_rate(self._calibration)
+        if rate is not None:
+            return rate
+        short = [
+            (d, res / n) for n, d, res in runtime + self._calibration if d >= n_features
+        ]
+        if not short:
+            return None
+        at = min(d for d, _ in short)
+        return max(r for d, r in short if d == at)
+
     def estimate(self, n_train: int, n_features: int) -> int | None:
         with self._lock:
             runtime = list(self._runtime)
-        slope = self._slope(runtime, n_features) or self._slope(
-            self._calibration, n_features
-        )
+        if self._width_dependent:
+            slope = self._width_dependent_slope(runtime, n_features)
+        else:
+            slope = self._slope(runtime, n_features) or self._slope(
+                self._calibration, n_features
+            )
         if slope is None:
             return None
         bounds = [
diff --git a/src/tabctx/serve/app.py b/src/tabctx/serve/app.py
index 476f219..4b59252 100644
--- a/src/tabctx/serve/app.py
+++ b/src/tabctx/serve/app.py
@@ -54,6 +54,7 @@
     scope_dataset_id,
 )
 from tabctx.serve.uploads import UploadStore
+from tabctx.serve.workers import WorkerDiedError
 
 logger = logging.getLogger("tabctx.serve")
 logging.basicConfig(level=logging.INFO)
@@ -575,7 +576,10 @@ def readyz(self):
             raise HTTPException(503, detail="; ".join(r for r in dead if r))
         # Worker processes measure their own device (this process holds no
         # CUDA context and must not create one on every GPU).
-        reports = [w.report() for w in self._workers]
+        try:
+            reports = [w.report() for w in self._workers]
+        except (WorkerDiedError, TimeoutError) as e:
+            raise HTTPException(503, detail=str(e)) from e
         real_gpu_memory = None
         per_device: dict[str, dict] = {}
         if reports:
diff --git a/src/tabctx/serve/workers.py b/src/tabctx/serve/workers.py
index 1add8e6..e44e51a 100644
--- a/src/tabctx/serve/workers.py
+++ b/src/tabctx/serve/workers.py
@@ -43,6 +43,7 @@
 import threading
 import time
 from concurrent.futures import Future, ThreadPoolExecutor
+from concurrent.futures import TimeoutError as FutureTimeoutError
 from dataclasses import dataclass
 from typing import Any
 
@@ -67,6 +68,9 @@
 _WORKER_THREADS = 16
 
 
+_NEVER_FITS = 1 << 62
+
+
 class WorkerDiedError(BackendComputeError):
     """A worker process exited with calls in flight."""
 
@@ -93,15 +97,22 @@ class _Failure:
 
 
 def _portable(error: BaseException) -> BaseException:
-    """The exception itself if it pickles (tabctx errors do), else a
-    RuntimeError carrying its type and message."""
+    """The exception itself if it survives a pickle round trip (tabctx
+    errors do), else a RuntimeError carrying its type and message. Checking
+    only dumps() isn't enough: an exception whose __init__ takes more than
+    its message pickles fine and then fails to unpickle in the reader."""
     try:
-        pickle.dumps(error)
+        pickle.loads(pickle.dumps(error))
         return error
     except Exception:  # noqa: BLE001
         return RuntimeError(f"{type(error).__name__}: {error}")
 
 
+# Bounded waits for calls that should answer promptly (fits and predicts
+# legitimately run for many minutes, so theirs are unbounded).
+QUICK_CALL_TIMEOUT_S = 30.0
+
+
 # ---- worker process --------------------------------------------------------
 
 
@@ -185,7 +196,12 @@ def handle(req_id: int, method: str, args: tuple, kwargs: dict) -> None:
         except BaseException as e:  # noqa: BLE001 -- returned to the caller
             reply = (req_id, False, _portable(e), snapshot())
         with send_lock:
-            conn.send(reply)
+            try:
+                conn.send(reply)
+            except Exception as e:  # noqa: BLE001 -- e.g. an unpicklable result
+                conn.send(
+                    (req_id, False, RuntimeError(f"unsendable reply: {e}"), snapshot())
+                )
 
     conn.send(_Ready(list(engine.backend_names), engine.default_backend, snapshot()))
     with ThreadPoolExecutor(_WORKER_THREADS, thread_name_prefix=f"w-{device}") as pool:
@@ -343,6 +359,13 @@ def _read_loop(self) -> None:
                 req_id, ok, payload, snap = self._conn.recv()
             except (EOFError, OSError):
                 break
+            except Exception:  # noqa: BLE001 -- an unreadable reply
+                # Its request can't be identified, so the stream is no longer
+                # trustworthy: retire the worker (its calls fail, health
+                # checks fail, the serving framework restarts the replica).
+                logger.exception("unreadable reply from worker for %s", self.device)
+                self._proc.kill()
+                break
             self._apply(snap)
             with self._state_lock:
                 fut = self._pending.pop(req_id, None)
@@ -375,7 +398,9 @@ def _apply(self, snap: Snapshot) -> None:
             if self._snapshot is None or snap.seq > self._snapshot.seq:
                 self._snapshot = snap
 
-    def _call(self, method: str, *args: Any, **kwargs: Any) -> Any:
+    def _call(
+        self, method: str, *args: Any, timeout_s: float | None = None, **kwargs: Any
+    ) -> Any:
         fut: Future = Future()
         with self._state_lock:
             if self._dead is not None:
@@ -385,11 +410,23 @@ def _call(self, method: str, *args: Any, **kwargs: Any) -> Any:
         try:
             with self._send_lock:
                 self._conn.send((req_id, method, args, kwargs))
-        except (OSError, ValueError) as e:
-            with self._state_lock:
-                self._pending.pop(req_id, None)
+        except (OSError, ValueError, EOFError) as e:
+            self._forget(req_id)
             raise WorkerDiedError(f"worker for {self.device} is gone: {e}") from e
-        return fut.result()
+        except Exception:  # unpicklable arguments: the caller's problem
+            self._forget(req_id)
+            raise
+        try:
+            return fut.result(timeout_s)
+        except FutureTimeoutError:
+            self._forget(req_id)
+            raise TimeoutError(
+                f"worker for {self.device} did not answer {method} in {timeout_s}s"
+            ) from None
+
+    def _forget(self, req_id: int) -> None:
+        with self._state_lock:
+            self._pending.pop(req_id, None)
 
     # ---- Shard: lock-free reads from the latest snapshot -------------
 
@@ -437,9 +474,20 @@ def admission_headroom_bytes(self, backend: str | None = None) -> int:
     def estimate_fit_bytes(
         self, n_train: int, n_features: int, backend: str | None = None
     ) -> int:
+        # A dead or unresponsive worker never fits: placement skips it rather
+        # than failing a fit other devices can take.
         if not self.alive:
-            return 1 << 62  # never fits: placement skips a dead worker
-        return self._call("estimate_fit_bytes", n_train, n_features, backend)
+            return _NEVER_FITS
+        try:
+            return self._call(
+                "estimate_fit_bytes",
+                n_train,
+                n_features,
+                backend,
+                timeout_s=QUICK_CALL_TIMEOUT_S,
+            )
+        except (WorkerDiedError, TimeoutError):
+            return _NEVER_FITS
 
     def fit(
         self,
@@ -473,4 +521,4 @@ def evict(self, dataset_id: str) -> None:
 
     def report(self) -> dict:
         """The worker's spill stats and real GPU memory (see _report)."""
-        return self._call("report")
+        return self._call("report", timeout_s=QUICK_CALL_TIMEOUT_S)
diff --git a/tests/integration/test_tabpfn_backend.py b/tests/integration/test_tabpfn_backend.py
index 43eb842..4e17dcc 100644
--- a/tests/integration/test_tabpfn_backend.py
+++ b/tests/integration/test_tabpfn_backend.py
@@ -64,7 +64,7 @@ def test_engine_stack_with_tabpfn(data, monkeypatch):
 
     from tabctx.backends.tabpfn import served_model_id
 
-    assert built.backend.name == served_model_id(settings.tabpfn.model_version.value)
+    assert built.default == served_model_id(settings.tabpfn.model_version.value)
 
     dataset_id = built.engine.fit(train_X, train_y, task="classification")
     outcome = built.engine.predict(dataset_id, test_X, return_proba=True)
diff --git a/tests/unit/test_engine_scheduling.py b/tests/unit/test_engine_scheduling.py
index a746df8..8d13eb1 100644
--- a/tests/unit/test_engine_scheduling.py
+++ b/tests/unit/test_engine_scheduling.py
@@ -201,3 +201,100 @@ def test_stats_never_wait_for_a_fit():
     assert engine.stats().n_cached_contexts == 0  # the cache lock is free
     backend.fit_gate.set()
     fitting.join(5)
+
+
+# ---- regressions from the review of 2e801ae --------------------------------
+
+
+def test_large_fit_and_large_predict_with_one_host_slot_do_not_deadlock():
+    """fit takes host slot -> _fit_lock; a large predict used to take
+    _fit_lock -> host slot. Forced: a small (slot-free) fit holds the
+    device while a large predict queues for it and then a large fit takes
+    the only host slot."""
+    from tabctx.memory.host import HostMemoryGuard
+
+    backend = GatedFake()
+    guard = HostMemoryGuard(48 * 1024**3, large_fit_cells=50)  # 1 slot
+    engine = TabctxEngine(
+        backend=backend,
+        estimator=CellEstimator(),
+        cache=ContextCacheManager(CAPACITY),
+        host_guard=guard,
+    )
+    engine.fit(*table(100, "x"), dataset_id="d")
+    backend.fit_gate = threading.Event()
+    small, _ = start(engine.fit, *table(10), dataset_id="small")  # 10 cells
+    assert backend.fit_entered.wait(5)
+    rows = [[0.0]] * (LIGHT_PREDICT_ROWS + 1)
+    predicting, pbox = start(engine.predict, "d", rows)
+    predicting.join(0.2)
+    fitting, fbox = start(engine.fit, *table(100), dataset_id="big")
+    fitting.join(0.2)
+    backend.fit_gate.set()
+    for t in (small, predicting, fitting):
+        t.join(10)
+    assert not predicting.is_alive() and not fitting.is_alive(), "deadlocked"
+    assert "error" not in pbox and "error" not in fbox
+
+
+def test_a_busy_context_is_evicted_when_nothing_else_makes_room(tmp_path):
+    """Back-to-back predicts kept a hot context pinned forever; the fit
+    went ahead without room and the cache went over capacity."""
+    engine, backend = build(tmp_path)
+    engine.fit(*table(600, "hot"), dataset_id="hot")
+    stop = threading.Event()
+    errors: list = []
+
+    def hammer():
+        while not stop.is_set():
+            try:
+                assert engine.predict("hot", [[0.0]]).predictions == ["hot"]
+            except Exception as e:  # noqa: BLE001
+                errors.append(e)
+                return
+
+    hammers = [threading.Thread(target=hammer, daemon=True) for _ in range(3)]
+    for t in hammers:
+        t.start()
+    engine.fit(*table(600), dataset_id="new")  # needs hot's room
+    stats = engine.stats()
+    assert stats.used_bytes <= stats.capacity_bytes
+    assert engine._cache.holds("new")
+    stop.set()
+    for t in hammers:
+        t.join(10)
+    assert errors == []  # predicts on "hot" restored it as needed
+
+
+def test_contexts_evicted_during_a_fit_are_freed_after_it_measures():
+    engine, backend = build()
+    engine.fit(*table(10), dataset_id="old")
+    backend.fit_gate = threading.Event()
+    fitting, _ = start(engine.fit, *table(100), dataset_id="big")
+    assert backend.fit_entered.wait(5)
+    engine.evict("old")
+    assert not engine._cache.holds("old")
+    assert [c.dataset_id for c in engine._parked] == ["old"]  # not freed yet
+    backend.fit_gate.set()
+    fitting.join(5)
+    assert engine._parked == []
+
+
+def test_a_restore_that_runs_out_of_memory_keeps_the_spilled_copy(tmp_path):
+    engine, backend = build(tmp_path)
+    engine.fit(*table(10, "x"), dataset_id="cold")
+    engine._cache.make_room(CAPACITY)
+    spill = engine._cache._spill
+    real_load, calls = spill.load, []
+
+    def oom_once(dataset_id):
+        calls.append(dataset_id)
+        if len(calls) == 1:
+            meta_path, payload_path = spill._paths(dataset_id)
+            assert meta_path.exists()
+            raise BackendComputeError("CUDA OOM while deserializing")
+        return real_load(dataset_id)
+
+    spill.load = oom_once
+    assert engine.predict("cold", [[0.0]]).predictions == ["x"]
+    assert len(calls) == 2  # retried; the file survived the first OOM
diff --git a/tests/unit/test_requirement.py b/tests/unit/test_requirement.py
index 4c10b7e..3ec1c13 100644
--- a/tests/unit/test_requirement.py
+++ b/tests/unit/test_requirement.py
@@ -278,3 +278,22 @@ def test_narrow_queries_use_their_own_width(self):
         assert self._model().estimate(150_000, 10) == pytest.approx(
             3 * GiB + rate * 50_000, rel=1e-6
         )
+
+
+def test_width_dependent_rate_ignores_a_short_runtime_fit_when_long_ones_exist():
+    """Review regression: one 500 x 100 runtime fit (fixed overhead / 500
+    rows = a huge per-row rate) displaced the calibration's long-fit rates
+    and doubled estimates for large tables."""
+    from tabctx.memory.requirement import ResidentRateModel
+
+    cal = (
+        FloorPoint(100_000, 10, GiB, resident_bytes=3 * GiB),
+        FloorPoint(100_000, 100, GiB, resident_bytes=int(3.1 * GiB)),
+    )
+    m = ResidentRateModel(cal, width_dependent=True)
+    before = (m.estimate(200_000, 10), m.estimate(200_000, 100))
+    m.record(500, 100, int(0.3 * GiB))
+    assert (m.estimate(200_000, 10), m.estimate(200_000, 100)) == before
+    # A long runtime fit does supersede the calibration's rate.
+    m.record(50_000, 100, 2 * GiB)
+    assert m.estimate(200_000, 100) != before[1]
diff --git a/tests/unit/test_spill.py b/tests/unit/test_spill.py
index 93ea73a..b057367 100644
--- a/tests/unit/test_spill.py
+++ b/tests/unit/test_spill.py
@@ -263,3 +263,22 @@ def test_failed_spill_is_a_real_miss_afterwards(self, tmp_path):
         cache.put(_ctx("bad", payload=lambda: None))
         cache.evict_one()
         assert not cache.has("bad")
+
+
+def test_put_keeps_a_replaced_id_visible_while_making_room(tmp_path):
+    """A same-device re-fit removed the old entry, then spilled others to
+    make room, then inserted: has() said 'nowhere' in between."""
+    store = GatedSpillStore(tmp_path)
+    store.gate.set()
+    cache = ContextCacheManager(200, spill_store=store)
+    cache.put(_ctx("refit", est_bytes=100))
+    cache.put(_ctx("other", est_bytes=100))
+    store.gate.clear()
+    store.entered.clear()
+    t = threading.Thread(target=cache.put, args=(_ctx("refit", est_bytes=200),))
+    t.start()
+    assert store.entered.wait(5)  # spilling "other" to make room
+    assert cache.has("refit")
+    store.gate.set()
+    t.join(5)
+    assert cache.holds("refit")
diff --git a/tests/unit/test_workers.py b/tests/unit/test_workers.py
index c1cacb6..f76488a 100644
--- a/tests/unit/test_workers.py
+++ b/tests/unit/test_workers.py
@@ -126,3 +126,15 @@ def test_startup_failure_is_raised_in_the_serving_process(tmp_path):
     blocker.write_text("")
     with pytest.raises(OSError):
         build_engine(_settings(blocker))
+
+
+def test_errors_that_cannot_round_trip_become_runtime_errors():
+    from tabctx.serve.workers import _portable
+
+    class TwoArgs(Exception):
+        def __init__(self, a, b):
+            super().__init__(f"{a}/{b}")
+
+    portable = _portable(TwoArgs(1, 2))
+    assert type(portable) is RuntimeError and "TwoArgs" in str(portable)
+    assert isinstance(_portable(DatasetNotFoundError("x")), DatasetNotFoundError)

From 6d88c4454e86ed0d947768761c93490e4305c09c Mon Sep 17 00:00:00 2001
From: Amrit Krishnan <amrit110@gmail.com>
Date: Sat, 17 Oct 2026 20:27:10 -0400
Subject: [PATCH 19/39] perf(tabicl): prepare a fit's CPU half before taking
 the device

TabICL's preprocessing (label encoding, the ensemble's per-feature power
transforms and shuffles) is 30-37% of a fit and needs no device memory.
The backend now splits fit() into prepare_fit (tabicl's fit with the kv
cache disabled) and fit_prepared (build the cache, measure, compress), and
the engine runs the first before taking the device, so a queued fit
preprocesses while another fit's GPU half runs. Measured on an A100: two
back-to-back fits 17-26% faster (2k x 20 to 40k x 100), the GPU half not
slowed by the concurrent preprocessing (7.75 s vs 7.94 s). Infeasible fits
are still rejected before any work.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/tabctx/backends/tabicl.py          | 51 ++++++++++++++++++++--
 src/tabctx/engine.py                   | 60 +++++++++++++++-----------
 tests/unit/test_engine_scheduling.py   | 51 ++++++++++++++++++++++
 tests/unit/test_tabicl_host_release.py |  7 +++
 4 files changed, 140 insertions(+), 29 deletions(-)

diff --git a/src/tabctx/backends/tabicl.py b/src/tabctx/backends/tabicl.py
index 93f2bff..1a086b5 100644
--- a/src/tabctx/backends/tabicl.py
+++ b/src/tabctx/backends/tabicl.py
@@ -51,6 +51,7 @@
 
 import logging
 import threading
+from dataclasses import dataclass
 from typing import Any
 
 import numpy as np
@@ -96,6 +97,26 @@ def _load_model(self: Any) -> None:
     cls._tabctx_restore_hook = True
 
 
+@dataclass
+class _PreparedFit:
+    """A fit whose CPU half is done (see TabICLBackend.prepare_fit)."""
+
+    task: Task
+    model: Any
+
+
+def _build_cache(model: Any) -> None:
+    """tabicl's fit() from the kv-cache step on (it follows the ensemble's
+    preprocessing in both TabICLClassifier and TabICLRegressor): pick the
+    cache mode from the estimator's kv_cache setting and build it."""
+    mode = model.kv_cache
+    if not mode:
+        model.model_kv_cache_ = None
+        return
+    model.cache_mode_ = "repr" if mode == "repr" else "kv"
+    model._build_kv_cache()
+
+
 def make_cache_thread_local(module: Any) -> None:
     """Give a shared TabICL backbone one active-context cache per thread.
 
@@ -237,13 +258,35 @@ def _stash_backbone(self, task: Task, model: Any) -> None:
         self._shared_backbones[task] = shared
 
     def fit(self, X: ArrayLike, y: ArrayLike, task: Task) -> Any:
-        import torch
-
+        return self.fit_prepared(self.prepare_fit(X, y, task))
+
+    def prepare_fit(self, X: ArrayLike, y: ArrayLike, task: Task) -> _PreparedFit:
+        """The CPU half of a fit: validation, label encoding and the
+        ensemble's preprocessing (per-feature power transforms, shuffles) --
+        tabicl's fit() up to building the kv cache, run with the cache
+        disabled. 30-37% of a fit's time (A100, 2k x 20 to 40k x 100); it
+        needs no device memory, so the engine runs it before taking the
+        device, overlapping another fit's GPU half."""
         # Fresh estimator instance every call (load-bearing -- see
         # backends/base.py), but the underlying pretrained nn.Module is
         # shared across fits after the first (see __init__).
         model = self._make_estimator(task)
+        kv_cache = model.kv_cache
+        model.kv_cache = False
+        try:
+            with device_scope(self._device):
+                model.fit(X, y)
+        finally:
+            model.kv_cache = kv_cache
+        return _PreparedFit(task, model)
 
+    def fit_prepared(self, prepared: _PreparedFit) -> Any:
+        """The GPU half: build the context's kv cache from a prepared fit,
+        measuring what it occupies. Re-runnable after a device OOM (tabicl
+        rebuilds the cache from scratch)."""
+        import torch
+
+        model = prepared.model
         device = self._device
         measuring = is_cuda(device)
         before_bytes = 0
@@ -256,7 +299,7 @@ def fit(self, X: ArrayLike, y: ArrayLike, task: Task) -> Any:
         self._last_fit_resident_bytes = None
         try:
             with device_scope(device):
-                model.fit(X, y)
+                _build_cache(model)
                 # The context as the fit built it, at full precision: what
                 # the fit needed resident, and so what admission must
                 # learn -- compression only shrinks it afterwards.
@@ -288,7 +331,7 @@ def fit(self, X: ArrayLike, y: ArrayLike, task: Task) -> Any:
             self._last_context_bytes = max(0, after_bytes - before_bytes)
             self._last_fit_resident_bytes = max(0, fit_resident or 0)
             self._last_fit_peak_bytes = max(0, peak_bytes - before_bytes)
-        self._stash_backbone(task, model)
+        self._stash_backbone(prepared.task, model)
         return model
 
     def _unpin_shared_module(self, model: Any) -> None:
diff --git a/src/tabctx/engine.py b/src/tabctx/engine.py
index 5f66bb4..0b5a650 100644
--- a/src/tabctx/engine.py
+++ b/src/tabctx/engine.py
@@ -312,14 +312,22 @@ def fit(
         estimator = self._estimators[backend_name]
 
         resolved_id = dataset_id or str(uuid.uuid4())
+        estimated = estimator.estimate_bytes(n_train, 0, n_features)
+        self._check_feasible(estimator, estimated, n_train, n_features)
+        # The fit's CPU half (backends that split it: TabICL's preprocessing,
+        # ~1/3 of a fit) runs before taking the device, overlapping the GPU
+        # half of whatever fit holds it now.
+        prepare = getattr(chosen, "prepare_fit", None)
+        prepared = prepare(X, y, task) if callable(prepare) else None
         # Lock order everywhere: host slot, _fit_lock, _predict_lock. Host
         # memory first, so a large fit waiting for a host slot doesn't hold
         # this device meanwhile.
         with self._host_slot(chosen, n_train, n_features), self._fit_lock:
-            estimated = estimator.estimate_bytes(n_train, 0, n_features)
-            self._admit(estimator, estimated, n_train, n_features)
+            self._make_room(estimator, estimated)
             with self._measuring():
-                payload = self._fit_backend(chosen, X, y, task, estimator, estimated)
+                payload = self._fit_backend(
+                    chosen, X, y, task, estimator, estimated, prepared
+                )
                 est_bytes = chosen.context_bytes_hint(n_train, n_features)
                 if est_bytes is None:
                     est_bytes = estimated
@@ -338,29 +346,24 @@ def fit(
             self._cache.put(context)
         return resolved_id
 
-    def _admit(
+    def _check_feasible(
         self,
         estimator: MemoryEstimator,
         estimated: int,
         n_train: int,
         n_features: int,
     ) -> None:
-        """Make room for a fit's estimated need, or reject it. Caller holds
-        _fit_lock, so no other residency growth races this check.
+        """Reject a fit that can never be admitted, before any work.
 
         Usage-aware admission (v0.9.0): what OOMs is the fit's need on top
-        of everything resident, so it must fit the headroom given current
-        cache usage. Evict-ahead-of-fit: when the cache is warm, cold
-        contexts are evicted (spilled, when a spill tier is attached) until
-        the need fits -- without this, a filling replica would reject fits
-        that are perfectly safe once a cold context steps aside (proven on
-        real hardware: a fourth ~13GB-peak fit OOMed an A100 with three
-        ~8GB contexts resident)."""
-        # Feasibility FIRST, against an EMPTY cache's headroom: a fit that
-        # can never be admitted must be rejected before evicting anything --
-        # otherwise one oversized request drains every tenant's context on
-        # its way to a 413 (observed live: a 200k x 60 request spilled 30
-        # contexts and then failed).
+        of everything resident, so under _fit_lock the fit evicts ahead
+        (_make_room) until it fits -- without that, a filling replica would
+        reject fits that are perfectly safe once a cold context steps aside
+        (proven on real hardware: a fourth ~13GB-peak fit OOMed an A100 with
+        three ~8GB contexts resident). Feasibility comes FIRST, against an
+        EMPTY cache's headroom: otherwise one oversized request drains every
+        tenant's context on its way to a 413 (observed live: a 200k x 60
+        request spilled 30 contexts and then failed)."""
         max_headroom = self._headroom(estimator, 0)
         if estimated > max_headroom:
             raise AdmissionRejected(
@@ -370,7 +373,6 @@ def _admit(
                 f"{max_headroom} byte maximum "
                 "admission headroom (even with an empty cache)"
             )
-        self._make_room(estimator, estimated)
 
     def _make_room(self, estimator: MemoryEstimator, needed: int) -> None:
         """Evict until `needed` bytes of headroom are free, or nothing is
@@ -423,20 +425,28 @@ def _fit_backend(
         task: Task,
         estimator: MemoryEstimator,
         estimated: int,
+        prepared: object | None = None,
     ) -> object:
-        """Fit, retrying once if it ran out of device memory while sharing
-        the device with a predict: after the failed attempt's tensors are
-        freed (outside the except block, whose traceback would keep them
-        alive), with predicts held off and room made again."""
-        try:
+        """Fit (its GPU half, when prepared), retrying once if it ran out of
+        device memory while sharing the device with a predict: after the
+        failed attempt's tensors are freed (outside the except block, whose
+        traceback would keep them alive), with predicts held off and room
+        made again."""
+
+        def attempt() -> object:
+            if prepared is not None:
+                return backend.fit_prepared(prepared)  # type: ignore[attr-defined]
             return backend.fit(X, y, task)
+
+        try:
+            return attempt()
         except BackendComputeError:
             pass
         gc.collect()
         release_cached_device_memory()
         with self._predict_lock:
             self._make_room(estimator, estimated)
-            return backend.fit(X, y, task)
+            return attempt()
 
     def _learn(
         self,
diff --git a/tests/unit/test_engine_scheduling.py b/tests/unit/test_engine_scheduling.py
index 8d13eb1..7a0d179 100644
--- a/tests/unit/test_engine_scheduling.py
+++ b/tests/unit/test_engine_scheduling.py
@@ -298,3 +298,54 @@ def oom_once(dataset_id):
     spill.load = oom_once
     assert engine.predict("cold", [[0.0]]).predictions == ["x"]
     assert len(calls) == 2  # retried; the file survived the first OOM
+
+
+class PreparingFake(GatedFake):
+    """A backend that splits its fit into a CPU prepare and a GPU half."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.prepared_during: list[bool] = []
+        self.engine: TabctxEngine | None = None
+
+    def prepare_fit(self, X, y, task):
+        # Record whether some fit held the device while this prepared.
+        self.prepared_during.append(self.engine._fit_lock.locked())
+        return (X, y, task)
+
+    def fit_prepared(self, prepared):
+        return GatedFake.fit(self, *prepared)
+
+
+def test_a_fit_prepares_while_another_fit_holds_the_device():
+    backend = PreparingFake()
+    engine = TabctxEngine(
+        backend=backend, estimator=CellEstimator(), cache=ContextCacheManager(CAPACITY)
+    )
+    backend.engine = engine
+    backend.fit_gate = threading.Event()
+    first, _ = start(engine.fit, *table(100), dataset_id="a")
+    assert backend.fit_entered.wait(5)
+    second, _ = start(engine.fit, *table(100), dataset_id="b")
+    for _ in range(200):
+        if len(backend.prepared_during) == 2:
+            break
+        threading.Event().wait(0.01)
+    assert backend.prepared_during == [False, True]  # b prepared during a's fit
+    backend.fit_gate.set()
+    first.join(5)
+    second.join(5)
+    assert engine._cache.holds("a") and engine._cache.holds("b")
+
+
+def test_an_infeasible_fit_is_rejected_before_preparing():
+    from tabctx.errors import AdmissionRejected
+
+    backend = PreparingFake()
+    engine = TabctxEngine(
+        backend=backend, estimator=CellEstimator(), cache=ContextCacheManager(CAPACITY)
+    )
+    backend.engine = engine
+    with pytest.raises(AdmissionRejected):
+        engine.fit(*table(2000), dataset_id="huge")  # 20,000 > 10,000 bytes
+    assert backend.prepared_during == []
diff --git a/tests/unit/test_tabicl_host_release.py b/tests/unit/test_tabicl_host_release.py
index ab5f041..09b5445 100644
--- a/tests/unit/test_tabicl_host_release.py
+++ b/tests/unit/test_tabicl_host_release.py
@@ -34,11 +34,18 @@ def clear_cache(self):
 
 
 class _Estimator:
+    """tabicl's two halves: fit() preprocesses (kv_cache off, as
+    prepare_fit runs it), _build_kv_cache() is the device work."""
+
     def __init__(self, fail: bool):
         self.model_ = _Module()
         self.fail = fail
+        self.kv_cache = "kv"
 
     def fit(self, X, y):
+        assert self.kv_cache is False  # prepare_fit disables the cache
+
+    def _build_kv_cache(self):
         if self.fail:
             raise RuntimeError("Failed to execute even with minimum batch size 1")
 

From 04d37515d44258bb38b8a627a76b2de817376b61 Mon Sep 17 00:00:00 2001
From: Amrit Krishnan <amrit110@gmail.com>
Date: Sat, 17 Oct 2026 20:28:52 -0400
Subject: [PATCH 20/39] fix(workers): close worker shards at interpreter exit
 (no spurious death reports)

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/tabctx/serve/workers.py | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/src/tabctx/serve/workers.py b/src/tabctx/serve/workers.py
index e44e51a..d13f335 100644
--- a/src/tabctx/serve/workers.py
+++ b/src/tabctx/serve/workers.py
@@ -37,11 +37,13 @@
 
 from __future__ import annotations
 
+import atexit
 import itertools
 import logging
 import pickle
 import threading
 import time
+import weakref
 from concurrent.futures import Future, ThreadPoolExecutor
 from concurrent.futures import TimeoutError as FutureTimeoutError
 from dataclasses import dataclass
@@ -71,6 +73,18 @@
 _NEVER_FITS = 1 << 62
 
 
+# Live shards, closed at interpreter exit -- before multiprocessing's own
+# exit handler (registered earlier, so run later) terminates the daemon
+# workers, which each reader thread would otherwise report as a death.
+_LIVE_SHARDS: weakref.WeakSet = weakref.WeakSet()
+
+
+@atexit.register
+def _close_live_shards() -> None:
+    for shard in list(_LIVE_SHARDS):
+        shard.close()
+
+
 class WorkerDiedError(BackendComputeError):
     """A worker process exited with calls in flight."""
 
@@ -255,6 +269,7 @@ def __init__(
         )
         self._proc.start()
         child.close()
+        _LIVE_SHARDS.add(self)
         self._send_lock = threading.Lock()
         self._state_lock = threading.Lock()
         self._pending: dict[int, Future] = {}

From dbdaf4164859d042ef1f5abeab600ce193be3917 Mon Sep 17 00:00:00 2001
From: Amrit Krishnan <amrit110@gmail.com>
Date: Sat, 17 Oct 2026 20:37:46 -0400
Subject: [PATCH 21/39] perf(tabicl): cache ssmax log(n) tensors (one GPU sync
 per length, not per call)

tabicl builds a 0-dim CUDA tensor from a Python float on every attention
call (ssmax._logn), a blocking host-to-device copy -- 42 syncs in a
40k x 100 fit, so the CPU never runs ahead of the GPU. Cached per
(n, device, dtype): bit-identical (max probability difference 0), fits
1.98 s -> 1.66 s at 10k x 50 and 11.1 -> 11.0 s at 40k x 100 on an A100.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/tabctx/backends/tabicl.py                | 32 ++++++++++++++++++++
 tests/unit/test_tabicl_thread_local_cache.py | 16 ++++++++++
 2 files changed, 48 insertions(+)

diff --git a/src/tabctx/backends/tabicl.py b/src/tabctx/backends/tabicl.py
index 1a086b5..ea27003 100644
--- a/src/tabctx/backends/tabicl.py
+++ b/src/tabctx/backends/tabicl.py
@@ -117,6 +117,36 @@ def _build_cache(model: Any) -> None:
     model._build_kv_cache()
 
 
+def install_cached_logn() -> bool:
+    """Replace tabicl's ssmax._logn with a cached equivalent.
+
+    _logn builds a 0-dim CUDA tensor from a Python float on every attention
+    call -- a host-to-device copy that blocks the CPU until the GPU queue
+    drains, 42 times in a 40k x 100 fit, so kernel launches never run ahead
+    of the GPU. The values depend only on (n, device, dtype); caching them
+    gives bit-identical results with one sync per distinct n. Measured on
+    an A100: 10,000 x 50 fits 1.98 s -> 1.66 s, 40,000 x 100 11.1 s -> 11.0 s,
+    max probability difference 0. Returns False (and changes nothing) if
+    tabicl no longer has the function."""
+    import functools
+
+    try:
+        from tabicl._model import ssmax
+    except ImportError:
+        return False
+    original = getattr(ssmax, "_logn", None)
+    if original is None or getattr(original, "_tabctx_cached", False):
+        return original is not None
+
+    @functools.lru_cache(maxsize=4096)
+    def cached(n: int, device: Any, dtype: Any) -> Any:
+        return original(n, device, dtype)
+
+    cached._tabctx_cached = True  # type: ignore[attr-defined]
+    ssmax._logn = cached
+    return True
+
+
 def make_cache_thread_local(module: Any) -> None:
     """Give a shared TabICL backbone one active-context cache per thread.
 
@@ -214,6 +244,8 @@ def __init__(
 
         _install_restore_hook(TabICLClassifier)
         _install_restore_hook(TabICLRegressor)
+        if not install_cached_logn():
+            logger.warning("tabicl.ssmax._logn not found: per-call GPU syncs remain")
 
     def _make_estimator(self, task: Task) -> Any:
         options = {
diff --git a/tests/unit/test_tabicl_thread_local_cache.py b/tests/unit/test_tabicl_thread_local_cache.py
index 0f9f42d..9e88cf1 100644
--- a/tests/unit/test_tabicl_thread_local_cache.py
+++ b/tests/unit/test_tabicl_thread_local_cache.py
@@ -58,3 +58,19 @@ def test_module_still_works_and_is_idempotent():
     assert m(torch.zeros(1, 2)).shape == (1, 2)
     assert "linear" in dict(m.named_modules())
     assert len(list(m.parameters())) == 2
+
+
+def test_cached_logn_matches_and_is_installed_once():
+    ssmax = pytest.importorskip("tabicl._model.ssmax")
+    from tabctx.backends.tabicl import install_cached_logn
+
+    original = ssmax._logn
+    try:
+        assert install_cached_logn()
+        patched = ssmax._logn
+        assert install_cached_logn() and ssmax._logn is patched  # idempotent
+        for n in (1, 7, 40_000):
+            expected = original(n, torch.device("cpu"), torch.float16)
+            assert torch.equal(patched(n, torch.device("cpu"), torch.float16), expected)
+    finally:
+        ssmax._logn = original

From 78ef4f2a2d92f5b3aada6c754fd8ad8aaa86c09c Mon Sep 17 00:00:00 2001
From: Amrit Krishnan <amrit110@gmail.com>
Date: Sat, 17 Oct 2026 20:52:40 -0400
Subject: [PATCH 22/39] feat(memory): extend the TabPFN-3.5 requirement grid to
 large tables

Six more measured shapes (A100, --max-slowdown 2): 200,000 x 10 needs
7.2 GiB free (5.8 resident), 100,000 x 100 3.6 GiB, 60,000 x 500 3.4 GiB,
30,000 x 1,000 2.8 GiB. Admission (L4 / A40) now reaches 200k / 200k rows
at 10 features, 100k / 100k at 100, 60k / 100k at 500 and 45k / 80k at
1,000 -- previously 45k / 5k / 1k / 500 on an L4.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../baselines/minfree-tabpfn-kv-a100.jsonl    |  6 +++
 .../calibration_tabpfn_a100_requirement.py    | 42 +++++++++++++++++++
 2 files changed, 48 insertions(+)

diff --git a/benchmarks/baselines/minfree-tabpfn-kv-a100.jsonl b/benchmarks/baselines/minfree-tabpfn-kv-a100.jsonl
index 562e1fc..dee0748 100644
--- a/benchmarks/baselines/minfree-tabpfn-kv-a100.jsonl
+++ b/benchmarks/baselines/minfree-tabpfn-kv-a100.jsonl
@@ -12,3 +12,9 @@
 {"rows": 500, "features": 1000, "status": "ok", "resident_gib": 1.18, "unconstrained_peak_gib": 7.31, "unconstrained_fit_s": 2.42, "min_free_gib": 5.7, "min_working_gib": 4.52, "fit_s_at_min": 3.863, "max_slowdown": 2.0}
 {"rows": 1000, "features": 1000, "status": "ok", "resident_gib": 1.22, "unconstrained_peak_gib": 7.42, "unconstrained_fit_s": 3.302, "min_free_gib": 5.73, "min_working_gib": 4.51, "fit_s_at_min": 4.81, "max_slowdown": 2.0}
 {"rows": 2000, "features": 1000, "status": "ok", "resident_gib": 1.32, "unconstrained_peak_gib": 7.64, "unconstrained_fit_s": 5.257, "min_free_gib": 5.82, "min_working_gib": 4.5, "fit_s_at_min": 6.9, "max_slowdown": 2.0}
+{"rows": 100000, "features": 100, "status": "ok", "resident_gib": 3.09, "unconstrained_peak_gib": 6.94, "unconstrained_fit_s": 83.702, "min_free_gib": 3.62, "min_working_gib": 0.53, "fit_s_at_min": 84.091, "max_slowdown": 2.0}
+{"rows": 30000, "features": 500, "status": "ok", "resident_gib": 1.96, "unconstrained_peak_gib": 6.65, "unconstrained_fit_s": 52.225, "min_free_gib": 2.24, "min_working_gib": 0.28, "fit_s_at_min": 54.784, "max_slowdown": 2.0}
+{"rows": 10000, "features": 1000, "status": "ok", "resident_gib": 1.98, "unconstrained_peak_gib": 7.07, "unconstrained_fit_s": 26.707, "min_free_gib": 2.53, "min_working_gib": 0.55, "fit_s_at_min": 41.427, "max_slowdown": 2.0}
+{"rows": 200000, "features": 10, "status": "ok", "resident_gib": 5.75, "unconstrained_peak_gib": 15.27, "unconstrained_fit_s": 202.293, "min_free_gib": 7.23, "min_working_gib": 1.48, "fit_s_at_min": 209.459, "max_slowdown": 2.0}
+{"rows": 60000, "features": 500, "status": "ok", "resident_gib": 2.82, "unconstrained_peak_gib": 7.39, "unconstrained_fit_s": 109.403, "min_free_gib": 3.36, "min_working_gib": 0.54, "fit_s_at_min": 115.808, "max_slowdown": 2.0}
+{"rows": 30000, "features": 1000, "status": "ok", "resident_gib": 2.55, "unconstrained_peak_gib": 7.01, "unconstrained_fit_s": 79.403, "min_free_gib": 2.82, "min_working_gib": 0.27, "fit_s_at_min": 84.952, "max_slowdown": 2.0}
diff --git a/src/tabctx/memory/calibration_tabpfn_a100_requirement.py b/src/tabctx/memory/calibration_tabpfn_a100_requirement.py
index 5801c6d..ad9f2f2 100644
--- a/src/tabctx/memory/calibration_tabpfn_a100_requirement.py
+++ b/src/tabctx/memory/calibration_tabpfn_a100_requirement.py
@@ -29,6 +29,12 @@
         working_bytes=6066641305,
         resident_bytes=3103113871,
     ),  # resident 2.89 GiB, unconstrained peak 9.64 GiB, fit 55.163s -> 58.662s at the floor
+    FloorPoint(
+        n_train=200000,
+        n_features=10,
+        working_bytes=1589137899,
+        resident_bytes=6174015488,
+    ),  # resident 5.75 GiB, unconstrained peak 15.27 GiB, fit 202.293s -> 209.459s at the floor
     FloorPoint(
         n_train=2000, n_features=100, working_bytes=622770257, resident_bytes=236223201
     ),  # resident 0.22 GiB, unconstrained peak 6.67 GiB, fit 0.845s -> 1.068s at the floor
@@ -41,6 +47,12 @@
     FloorPoint(
         n_train=20000, n_features=100, working_bytes=311385128, resident_bytes=858993459
     ),  # resident 0.8 GiB, unconstrained peak 7.3 GiB, fit 9.669s -> 11.485s at the floor
+    FloorPoint(
+        n_train=100000,
+        n_features=100,
+        working_bytes=569083166,
+        resident_bytes=3317862236,
+    ),  # resident 3.09 GiB, unconstrained peak 6.94 GiB, fit 83.702s -> 84.091s at the floor
     FloorPoint(
         n_train=1000, n_features=500, working_bytes=2448131358, resident_bytes=869730877
     ),  # resident 0.81 GiB, unconstrained peak 7.0 GiB, fit 2.11s -> 3.492s at the floor
@@ -53,6 +65,18 @@
         working_bytes=2426656522,
         resident_bytes=1181116006,
     ),  # resident 1.1 GiB, unconstrained peak 7.38 GiB, fit 8.321s -> 10.432s at the floor
+    FloorPoint(
+        n_train=30000,
+        n_features=500,
+        working_bytes=300647710,
+        resident_bytes=2104533975,
+    ),  # resident 1.96 GiB, unconstrained peak 6.65 GiB, fit 52.225s -> 54.784s at the floor
+    FloorPoint(
+        n_train=60000,
+        n_features=500,
+        working_bytes=579820584,
+        resident_bytes=3027951943,
+    ),  # resident 2.82 GiB, unconstrained peak 7.39 GiB, fit 109.403s -> 115.808s at the floor
     FloorPoint(
         n_train=500,
         n_features=1000,
@@ -71,6 +95,18 @@
         working_bytes=4831838208,
         resident_bytes=1417339207,
     ),  # resident 1.32 GiB, unconstrained peak 7.64 GiB, fit 5.257s -> 6.9s at the floor
+    FloorPoint(
+        n_train=10000,
+        n_features=1000,
+        working_bytes=590558003,
+        resident_bytes=2126008811,
+    ),  # resident 1.98 GiB, unconstrained peak 7.07 GiB, fit 26.707s -> 41.427s at the floor
+    FloorPoint(
+        n_train=30000,
+        n_features=1000,
+        working_bytes=289910292,
+        resident_bytes=2738041651,
+    ),  # resident 2.55 GiB, unconstrained peak 7.01 GiB, fit 79.403s -> 84.952s at the floor
 )
 
 A100_40GB_TABPFN_KV_REQUIREMENT_GRID: tuple[Observation, ...] = (
@@ -78,14 +114,20 @@
     Observation(n_train=10000, n_features=10, real_bytes=3124588707),
     Observation(n_train=50000, n_features=10, real_bytes=4574140170),
     Observation(n_train=100000, n_features=10, real_bytes=9169755176),
+    Observation(n_train=200000, n_features=10, real_bytes=7763153387),
     Observation(n_train=2000, n_features=100, real_bytes=858993459),
     Observation(n_train=5000, n_features=100, real_bytes=1299227607),
     Observation(n_train=10000, n_features=100, real_bytes=1170378588),
     Observation(n_train=20000, n_features=100, real_bytes=1170378588),
+    Observation(n_train=100000, n_features=100, real_bytes=3886945402),
     Observation(n_train=1000, n_features=500, real_bytes=3317862236),
     Observation(n_train=2000, n_features=500, real_bytes=3393024163),
     Observation(n_train=5000, n_features=500, real_bytes=3607772528),
+    Observation(n_train=30000, n_features=500, real_bytes=2405181685),
+    Observation(n_train=60000, n_features=500, real_bytes=3607772528),
     Observation(n_train=500, n_features=1000, real_bytes=6120328396),
     Observation(n_train=1000, n_features=1000, real_bytes=6152540651),
     Observation(n_train=2000, n_features=1000, real_bytes=6249177415),
+    Observation(n_train=10000, n_features=1000, real_bytes=2716566814),
+    Observation(n_train=30000, n_features=1000, real_bytes=3027951943),
 )

From 41c7b441d2a6d8563d2ac91074755c0d9bcbec8f Mon Sep 17 00:00:00 2001
From: Amrit Krishnan <amrit110@gmail.com>
Date: Sat, 17 Oct 2026 20:53:21 -0400
Subject: [PATCH 23/39] docs: roadmap and changelog for worker processes,
 scheduling, profiling, TabPFN calibration

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 CHANGELOG.md | 10 +++++++++
 ROADMAP.md   | 63 ++++++++++++++++++++++++++++++++++++++--------------
 2 files changed, 56 insertions(+), 17 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index b1a5637..b5b0d24 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,16 +5,26 @@
 ## [0.12.0] - unreleased
 ### Added
 - **Multi-GPU replicas** (#30). `MultiDeviceEngine` runs one `TabctxEngine` per device -- backends bound to `cuda:i`, its own cache, lock and spill subdirectory -- behind the single-engine API, built automatically when a replica sees several GPUs (`TABCTX_DEVICES=auto`, or an explicit list). Placement is aimed at the expected workload, a few concurrent users with small and large tables: a fit goes where its estimated peak fits without eviction, on the least active device (in-flight calls, then calls in the last 30s), else to the device with the most headroom. Predicts route to the device holding the context (resident first, then spilled), found from lock-free residency/spill snapshots, so routing never waits on another GPU's work and fits and predicts on different devices run concurrently. Per-dataset_id locks serialize fits and evicts of one id (no duplicate copies from retried fits); a re-fit lands on its device before other copies are removed, so a failed re-fit leaves the previous context serving. Unknown models are rejected (422) before anything changes. `TABCTX_SPILL_CAPACITY_BYTES` is the replica's total, split across devices' spill subdirectories; a spill layout change (one vs several devices) is logged. `TABCTX_DEVICES` rejects mixed CPU/GPU lists, a bare `cuda` among indexed GPUs, and indices the process can't see. `/readyz` gains a per-device breakdown (contexts, bytes, activity, real GPU memory); `/v1/tabctx/limits` reports the best single device. A single GPU builds exactly the engine it did before.
+- **One worker process per GPU** (#39, `TABCTX_WORKER_PROCESSES`, default on for multi-GPU replicas). Fits are host-bound, so threads of one process driving several GPUs contend for the GIL: on 4x A100, four concurrent fits ran at 2.1-3.1x aggregate from threads and 3.96-4.3x from processes (`benchmarks/bench_scaling.py`). `serve/workers.py` runs each device's engine in its own process behind a `ProcessShard` proxy; the router reads versioned residency/headroom snapshots that every reply carries, so its decisions stay lock-free. A dead worker fails its calls, is skipped by placement, and fails `/readyz` and Ray Serve's `check_health` so the replica restarts. Nothing depends on Ray. Under Ray Serve on 4x A100 with 6 mixed users: 2.58 -> 5.0 predicts/s, small-fit p50 1.6 -> 0.6 s.
+- **Light predicts run beside a fit** on the same GPU (#39). A predict against a resident context with at most 2,048 rows no longer waits for a fit on its device (it waited up to 266 s on the L4 rehearsal): ~0.1 s p50 during a 40k x 100 fit on an A100, at ~15% fit slowdown. Restores and large predicts still wait; contexts in use are pinned against evict-ahead (and evicted with predicts held off when nothing else makes room); a fit that runs out of memory beside a predict retries alone; allocator measurements are taken with predicts quiesced, and memory freed during a fit is released only after it has measured itself.
+- **TabPFN-3.5 requirement calibration** (A100, 20 shapes up to 200k x 10 and 30k x 1,000). TabPFN also sizes its working memory from what is free (20k x 100 peaks at 7.3 GiB unconstrained but runs within 2x its time on 1.1 GiB), so admission bounds resident + working floor for it too, instead of TabICL's peak formula: up to 100k rows at 100 features and 80k at 1,000 on an A40 (was 5k and 500 on an L4). Its per-row resident rate grows with features, which `ResidentRateModel(width_dependent=True)` respects.
+- `benchmarks/bench_scaling.py` (threads vs processes across GPUs), `bench_overlap.py` (predict latency and correctness during a fit), `gke/serve_workload.sh` (Ray Serve + workload inside the devbox).
 - `benchmarks/bench_workload.py`: concurrent users looping fit-then-predict over a small/medium/large table mix (large ones via CSV upload), with per-class latency percentiles, 413s, errors and the per-device spread.
 - `benchmarks/sweep_admission.py`: the table-size sweep (rows up to the first 413 per feature count, fit and predict timed).
 - ROADMAP: multi-GPU and CUDA-level track (#30-#37), PyTorch-level first.
 
 ### Changed
+- **TabICL fits prepare their CPU half before taking the device**: preprocessing (~30-37% of a fit) overlaps another fit's GPU half -- back-to-back fits 17-26% faster on an A100.
+- **tabicl's per-call log(n) host-to-device syncs are cached** (bit-identical): 10k x 50 fits 1.98 -> 1.66 s on an A100.
+- The cache's spill reads and writes run outside its lock, so bookkeeping and light predicts never wait on multi-GB disk I/O.
 - **TabPFN contexts share the model: ~42x less GPU memory per context.** TabPFN builds a fresh model for every estimator unless its built-model cache is on, so every cached context held its own ~850 MB copy of the weights on the GPU (measured on an A100: 4 contexts of a 300x6 table took 3.4 GB). `TabPFNBackend` now enables that cache (`TABPFN_MODEL_CACHE_SIZE`, default 16 entries, keyed by device and task): one shared model per (device, task), about 1.7 GB fixed per GPU, and a context costs only its own state (~20 MB for that table). Interleaved contexts predict bit-identically to each context alone.
 - **TabPFN spill blobs no longer contain the model.** A spilled TabPFN payload was ~900 MB, mostly weights, and a restore built another private copy on the GPU. References to the shared model are now pickled as tokens and re-attached on load: a 2,000x20 context spills as 99 MB and restores with identical predictions and no extra GPU memory. Older blobs still load.
 - **Weights load at replica start** (`warmup()` on both backends, called by the factory) instead of inside the first tenant's fit, which reported the checkpoint as part of its context (a 200x6 TabPFN context was booked at 855 MB, TabICL's at 226 MB) and paid the load latency.
 
 ### Fixed
+- **Router 404 for a context being spilled** (4xL4 rehearsal: a fit returned 200 and the next predict 404'd 27 ms later): a context moving between the GPU cache and the spill tier was listed by neither; it is now tracked in transit, and predict re-locates once after a miss.
+- **TabICL shared-backbone cache was per-module, not per-call**: interleaved calls could read another context's cache. It is now thread-local.
+- A device OOM while restoring a spilled context no longer deletes the spilled copy.
 - **First fit on a GPU failed** with "Invalid device argument": device-explicit memory-stat calls need CUDA initialized, which the no-argument forms used to do implicitly. Backends now initialize CUDA at construction.
 - **Spill tier warm restart:** the index rebuilt at startup was keyed by the hex-encoded file name instead of the dataset_id, so lookups, explicit deletes and disk-pressure cleanup all missed contexts a previous process had spilled (restores still worked because `load()` reads by path).
 - **TabICL spill restore reloaded the checkpoint:** tabicl's `__setstate__` calls `_load_model()`, which loaded the full checkpoint (and a second backbone onto the GPU, unaccounted) on every restore before tabctx swapped in the shared one. Spilled TabICL payloads now carry a task tag and restore straight onto the shared backbone; untagged blobs from older versions still load.
diff --git a/ROADMAP.md b/ROADMAP.md
index 3a79a9f..de24322 100644
--- a/ROADMAP.md
+++ b/ROADMAP.md
@@ -56,24 +56,53 @@ the kernel work. Everything is developed and tested on a GKE GPU node
 first, then validated on prod's A40s (Ampere sm_86: no FlashAttention-3,
 no FP8).
 
-5. **Multi-GPU replica** ([#30](https://github.com/VectorInstitute/tabctx/issues/30), in progress): one engine per device,
+5. **Multi-GPU replica** ([#30](https://github.com/VectorInstitute/tabctx/issues/30), done -- PR #38/#40): one engine per device,
    fits placed where they fit without eviction on the least active
-   device, predicts routed to the owning device.
-6. **Profiling harness** ([#31](https://github.com/VectorInstitute/tabctx/issues/31)): per-stage time and peak-memory breakdown;
-   picks and sizes everything below. Hypothesis to confirm: the fit
-   peak is per-cell activations in the column/row stages, not attention.
-7. **Per-device calibration grids** ([#37](https://github.com/VectorInstitute/tabctx/issues/37)): admission on non-A100 cards
-   falls back to the conservative formula and is the first ceiling a
-   large-table user hits.
-8. **Tiled/fused column+row stage kernels** ([#32](https://github.com/VectorInstitute/tabctx/issues/32)): cut the fit peak, raising
-   the per-GPU table ceiling (FlashAttention-style tiling; chunked
-   prefill for the context KV).
-9. **Paged context KV + cross-context varlen predict** ([#33](https://github.com/VectorInstitute/tabctx/issues/33)): throughput with
-   many users (PagedAttention + continuous batching analogue).
-10. **INT8 context KV** ([#34](https://github.com/VectorInstitute/tabctx/issues/34)): about 2x contexts per GPU.
-11. **Context parallelism across GPUs** ([#35](https://github.com/VectorInstitute/tabctx/issues/35)): tables larger than one GPU
-    (ring attention over the replica's devices).
-12. **CUDA graphs for small predicts** ([#36](https://github.com/VectorInstitute/tabctx/issues/36)).
+   device, predicts routed to the owning device. **One worker process per
+   GPU** ([#39](https://github.com/VectorInstitute/tabctx/issues/39)): fits are host-bound (the driving thread is busy ~90% of
+   wall time), so threads of one process contend for the GIL -- 4x A100,
+   four concurrent fits: 2.1-3.1x aggregate from threads, 3.96-4.3x from
+   processes. Under Ray Serve, 6 mixed users: 2.58 -> 5.0 predicts/s.
+6. **Profiling** ([#31](https://github.com/VectorInstitute/tabctx/issues/31), first pass done): a 40k x 100 TabICL fit on an
+   A100 is ~30% CPU preprocessing (per-feature Yeo-Johnson power
+   transforms in scipy dominate) and ~70% GPU half, which is ~87% GPU-busy
+   (flash attention 35%, layer norm 21% -- memory-bound over
+   [4, 40, 40000, 128] activations -- copies 15%). Done from it: the CPU
+   half is prepared before taking the device (overlaps another fit's GPU
+   half: back-to-back fits 17-26% faster), and tabicl's per-call
+   log(n) host-to-device syncs are cached (bit-identical; 10k x 50 fits
+   16% faster). Next from it: fuse the column stage's layer norms
+   (`torch.compile`, #32/#36), and a GPU Yeo-Johnson fit for isolated-fit
+   latency (~3 ms vs 2.9 s at 40k x 100; not bit-identical to sklearn's
+   lambdas, so behind an accuracy study).
+7. **Per-device calibration grids** ([#37](https://github.com/VectorInstitute/tabctx/issues/37), TabPFN done): TabPFN-3.5 has its
+   own requirement grid (A100, 20 shapes): it, too, sizes working memory
+   from what is free, so admission reaches 100k rows at 100 features on an
+   A40 (was 5k on an L4 under TabICL's formula). Still to do: L4/A40-
+   measured grids (A100 grids scale by capacity today).
+8. **Head-of-line scheduling** ([#39](https://github.com/VectorInstitute/tabctx/issues/39), done): light predicts (resident
+   context, <= 2048 rows) run beside a fit on the same GPU instead of
+   waiting for it -- ~0.1 s p50 during a 40k x 100 fit instead of its
+   full 11 s, 0 wrong answers (thread-local backbone cache); in-use
+   contexts are pinned against evict-ahead; restores and large predicts
+   wait. Small-fit p95 under load 10.6 s -> 1.4 s.
+9. **Tiled/fused column+row stage kernels** ([#32](https://github.com/VectorInstitute/tabctx/issues/32)): cut the fit peak and the
+   memory-bound layer norms, raising the per-GPU table ceiling.
+10. **Paged context KV + cross-context varlen predict** ([#33](https://github.com/VectorInstitute/tabctx/issues/33)): throughput
+    with many users (PagedAttention + continuous batching analogue).
+11. **INT8 context KV** ([#34](https://github.com/VectorInstitute/tabctx/issues/34), done for TabICL): ~0.52x resident, identical
+    labels on the evaluation sets; TabPFN already defaults to int8.
+12. **Context parallelism across GPUs** ([#35](https://github.com/VectorInstitute/tabctx/issues/35)): tables larger than one GPU
+    (ring attention over the replica's devices). Disk offload already
+    lets one A100 fit 130k x 500 (65M cells) in 26 min.
+13. **CUDA graphs / torch.compile** ([#36](https://github.com/VectorInstitute/tabctx/issues/36)): the GPU half is now mostly GPU-
+    bound, so the remaining gains are fusion (layer norm, elementwise)
+    rather than launch overhead.
+14. **Serving-framework-agnostic core** ([#39](https://github.com/VectorInstitute/tabctx/issues/39)): the router, per-GPU workers
+    and scheduling are plain Python with no Ray dependency; a Dynamo
+    worker hosts the same core. Next: residency events and pluggable
+    placement so a platform scheduler can co-optimize LLMs and tabular
+    models.
 
 PyTorch-level first: tiling and chunking, `torch.compile` (generated Triton
 kernels, CUDA graphs via `reduce-overhead`), FlexAttention for paged and

From 4d3382fa064731b28740599382b80ab122062129 Mon Sep 17 00:00:00 2001
From: Amrit Krishnan <amrit110@gmail.com>
Date: Sat, 17 Oct 2026 20:53:39 -0400
Subject: [PATCH 24/39] docs(readme): worker processes and predicts beside fits

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 README.md | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/README.md b/README.md
index 99ab65e..8e3bb77 100644
--- a/README.md
+++ b/README.md
@@ -223,6 +223,16 @@ behind the same API. No client changes and no routing header needed.
   retried fit can't leave two copies. A re-fit lands first and only then
   replaces the old copy: a failed re-fit (413, OOM) leaves the previous
   context serving.
+- **A process per GPU.** Each GPU's engine runs in its own worker process
+  (`TABCTX_WORKER_PROCESSES=1`, the default with several GPUs; `0` runs
+  them as threads). Fits are host-bound, so threads would contend for
+  Python's GIL: on 4x A100, four concurrent fits ran at ~4x aggregate
+  from processes and 2-3x from threads. A worker that dies fails the
+  replica's health check, so the serving framework restarts it.
+- **Small predicts don't wait for fits.** A predict of up to 2,048 rows
+  against a resident context runs beside a fit on the same GPU (~0.1 s
+  instead of the fit's full duration); larger predicts and restores of
+  spilled contexts wait their turn.
 - **The replica must be given the GPUs.** Under Ray Serve that means
   `ray_actor_options.num_gpus: N`, or, for a `num_gpus: 0` actor on a
   GPU pod (as in inference-platform's prod config),

From 44471ff36a35a9258ed69f10d31d637d956c5e8b Mon Sep 17 00:00:00 2001
From: Amrit Krishnan <amrit110@gmail.com>
Date: Sat, 17 Oct 2026 21:37:44 -0400
Subject: [PATCH 25/39] feat(memory): admit host memory by estimated bytes, not
 large-fit slots

The host guard gave every call that might offload a whole 48 GiB slot
(the A100 peak of a 10M-cell TabICL fit): on a 4xL4 node 40k x 100 fits
-- which offload on a 24 GB card -- serialized on 2-3 slots across all
four GPUs (large fits p50 195 s on the rehearsal), and every TabPFN fit
over 3M cells took a slot though TabPFN never exceeds 6.7 GiB of host
memory (30M cells, bench_host_memory on an A100). Calls now reserve their
own estimated bytes against the budget: TabICL ~5 KB per cell when it may
offload, TabPFN 600 B per cell (worst measured ~530), else one large fit
above the cell threshold. A call bigger than the budget runs alone.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/tabctx/backends/tabicl.py  |   9 +++
 src/tabctx/backends/tabpfn.py  |   6 ++
 src/tabctx/engine.py           |  17 ++++--
 src/tabctx/memory/host.py      |  97 ++++++++++++++++--------------
 tests/unit/test_host_memory.py | 106 ++++++++++++++++++++++++---------
 5 files changed, 155 insertions(+), 80 deletions(-)

diff --git a/src/tabctx/backends/tabicl.py b/src/tabctx/backends/tabicl.py
index ea27003..a105a54 100644
--- a/src/tabctx/backends/tabicl.py
+++ b/src/tabctx/backends/tabicl.py
@@ -72,6 +72,7 @@
 _restoring = threading.local()
 _warned = [False]  # one missing-hook warning per process
 # Offloaded stage-output bytes per cell (see TabICLBackend.may_offload).
+_HOST_BYTES_PER_OFFLOADED_CELL = 5_000
 _OUTPUT_BYTES_PER_CELL = 2560
 logger = logging.getLogger(__name__)
 _TASK_TAGS = {"classification": b"C", "regression": b"R"}
@@ -396,6 +397,14 @@ def may_offload(self, n_rows: int, n_features: int) -> bool:
         free, _ = torch.cuda.mem_get_info(torch.device(self._device))
         return _OUTPUT_BYTES_PER_CELL * n_rows * n_features > 0.4 * free
 
+    def host_bytes_estimate(self, n_rows: int, n_features: int) -> int:
+        """Host memory a call over n_rows x n_features may use: its offloaded
+        stage outputs when it may offload (measured ~46 GiB for 10M-cell
+        fits on an A100, ~4.9 KB per cell), else nothing worth reserving."""
+        if not self.may_offload(n_rows, n_features):
+            return 0
+        return _HOST_BYTES_PER_OFFLOADED_CELL * n_rows * n_features
+
     def _release_host_memory(self, model: Any) -> None:
         """Return the offload staging memory a large call leaves pinned.
 
diff --git a/src/tabctx/backends/tabpfn.py b/src/tabctx/backends/tabpfn.py
index 8deadb4..f2d2e4a 100644
--- a/src/tabctx/backends/tabpfn.py
+++ b/src/tabctx/backends/tabpfn.py
@@ -162,6 +162,12 @@ def fit(self, X: ArrayLike, y: ArrayLike, task: Task) -> Any:
             self._shared_models.setdefault((task, index), module)
         return model
 
+    def host_bytes_estimate(self, n_rows: int, n_features: int) -> int:
+        """Host memory a fit or predict may use (preprocessing copies):
+        measured on an A100 pod at most ~530 B per cell (100,000 x 100 grew
+        RSS 4.9 GiB; 30M-cell tables 6-6.7 GiB), budgeted at 600."""
+        return 600 * n_rows * n_features
+
     def predict(
         self, payload: Any, X_test: ArrayLike, return_proba: bool = False
     ) -> PredictOutcome:
diff --git a/src/tabctx/engine.py b/src/tabctx/engine.py
index 0b5a650..cbdb175 100644
--- a/src/tabctx/engine.py
+++ b/src/tabctx/engine.py
@@ -479,18 +479,23 @@ def _learn(
             )
 
     def _host_slot(self, backend: TabularICLBackend, n_rows: int, n_features: int):
-        """A process-wide large-call slot when this call may offload to host
-        memory (the backend's may_offload, else the guard's cell threshold);
-        a no-op without a guard."""
-        if self._host_guard is None:
+        """Reserve this call's estimated host memory against the process-
+        wide (or node-wide) budget: the backend's host_bytes_estimate, else
+        one large fit's worth when it may offload (may_offload, or the
+        guard's cell threshold). A no-op without a guard."""
+        guard = self._host_guard
+        if guard is None:
             return contextlib.nullcontext()
+        estimate = getattr(backend, "host_bytes_estimate", None)
+        if callable(estimate):
+            return guard.reserve(estimate(n_rows, n_features))
         may_offload = getattr(backend, "may_offload", None)
         large = (
             may_offload(n_rows, n_features)
             if callable(may_offload)
-            else self._host_guard.is_large(n_rows, n_features)
+            else guard.is_large(n_rows, n_features)
         )
-        return self._host_guard.reserve(large)
+        return guard.reserve(guard.large_fit_bytes if large else 0)
 
     def _restore(self, dataset_id: str) -> None:
         """Bring a spilled context back onto the device (DatasetNotFoundError
diff --git a/src/tabctx/memory/host.py b/src/tabctx/memory/host.py
index 3aefaac..302eff7 100644
--- a/src/tabctx/memory/host.py
+++ b/src/tabctx/memory/host.py
@@ -1,23 +1,26 @@
 """Host (CPU) memory admission: the limit GPU admission doesn't see.
 
-Large TabICL fits offload stage outputs to pinned host memory while they
-run. Measured on an A100 pod (benchmarks/bench_host_memory.py,
+Large TabICL fits offload stage outputs to host memory while they run.
+Measured on an A100 pod (benchmarks/bench_host_memory.py,
 benchmarks/baselines/host-tabicl-a100.jsonl): fits up to ~5M cells grow
-host RSS by < 0.5 GiB, but 10-15M-cell fits (20,000 x 500, 100,000 x 100)
-peak at ~46 GiB of host memory -- transient, returned afterwards (see
-TabICLBackend._release_host_memory). One such fit fits a 64 GiB pod; two
-at once, on two GPUs of one replica, get it OOM-killed with every tenant's
-context.
-
-HostMemoryGuard lets at most budget // large_fit_bytes host-offloading
-calls run at once, process-wide (shared by every device's engine). Which
-calls may offload is the backend's call (TabICLBackend.may_offload mirrors
-tabicl's trigger: a stage output above half the device's currently free
-GPU memory); backends without that hint use a cell-count threshold. A
-large call waits for a slot instead of exhausting the pod; small ones
-never wait. Each slot's call should budget at most 1/slots of the host
-headroom for itself (see TabICLBackend's cpu_safety_factor), so
-concurrent calls can't all claim the same free memory.
+host RSS by < 0.5 GiB, but 10M-cell fits that offload (20,000 x 500,
+100,000 x 100) peak at ~46 GiB -- transient, returned afterwards (see
+TabICLBackend._release_host_memory). TabPFN stays small: at most 6.7 GiB
+at 30M cells (host-tabpfn-a100.jsonl). One huge fit fits a 64 GiB pod;
+two at once get it OOM-killed with every tenant's context.
+
+HostMemoryGuard admits calls by their estimated host bytes against one
+budget for the whole process -- or, built with mp_context, for every
+per-GPU worker process at once (host RAM is shared by the node). Each
+backend estimates its own calls (host_bytes_estimate: TabICL ~5 KB per
+cell when it may offload, TabPFN ~0.6 KB per cell); backends without an
+estimate count as one large fit above a cell threshold. A call waits
+until its bytes fit beside those already admitted (a call bigger than the
+whole budget runs alone, as the only one); calls estimated at 0 never
+wait. Calls that do offload budget their own staging against ``slots``,
+the number of large fits the budget holds (see TabICLBackend's
+cpu_safety_factor), so concurrent ones can't all claim the same free
+memory.
 """
 
 from __future__ import annotations
@@ -31,7 +34,8 @@
 
 GiB = 1024**3
 logger = logging.getLogger(__name__)
-# Calibrated on an A100-40GB (see module docstring).
+# Calibrated on an A100-40GB (see module docstring): the fallback for
+# backends without a host_bytes_estimate.
 LARGE_FIT_CELLS = 3_000_000
 LARGE_FIT_HOST_BYTES = 48 * GiB
 
@@ -88,11 +92,6 @@ def _read_limit(limit_paths: tuple[str, ...], meminfo_path: str) -> int | None:
 
 
 class HostMemoryGuard:
-    """Process-wide by default. With mp_context, its semaphore and counters
-    are multiprocessing primitives, so one guard passed to per-GPU worker
-    processes (as a Process argument) bounds them all: host RAM is shared
-    by every process on the node."""
-
     def __init__(
         self,
         budget_bytes: int,
@@ -100,6 +99,8 @@ def __init__(
         large_fit_bytes: int = LARGE_FIT_HOST_BYTES,
         mp_context: Any | None = None,
     ) -> None:
+        self.budget_bytes = budget_bytes
+        self.large_fit_bytes = large_fit_bytes
         self.slots = max(1, budget_bytes // large_fit_bytes)
         if budget_bytes < large_fit_bytes:
             logger.warning(
@@ -111,42 +112,46 @@ def __init__(
             )
         self._large_fit_cells = large_fit_cells
         if mp_context is None:
-            self._semaphore: Any = threading.BoundedSemaphore(self.slots)
-            self._lock: Any = threading.Lock()
-            self._counts: Any = [0, 0]  # waiting, active
+            self._cond: Any = threading.Condition()
+            self._counts: Any = [0, 0, 0]  # reserved bytes, active, waiting
         else:
-            self._semaphore = mp_context.BoundedSemaphore(self.slots)
-            self._lock = mp_context.Lock()
-            self._counts = mp_context.Array("i", 2, lock=False)
+            self._cond = mp_context.Condition(mp_context.Lock())
+            self._counts = mp_context.Array("q", 3, lock=False)
 
     def is_large(self, n_rows: int, n_features: int) -> bool:
-        """The fallback test for backends that can't say whether a call may
-        offload: a cell count below which no offload was ever measured."""
+        """The fallback test for backends that can't estimate their host
+        use: a cell count below which no offload was ever measured."""
         return n_rows * n_features >= self._large_fit_cells
 
     @contextlib.contextmanager
-    def reserve(self, large: bool) -> Iterator[None]:
-        """Hold a large-call slot for the duration (no-op when not large)."""
-        if not large:
+    def reserve(self, nbytes: int) -> Iterator[None]:
+        """Hold `nbytes` of the host budget for the duration (no-op for 0)."""
+        if nbytes <= 0:
             yield
             return
-        with self._lock:
-            self._counts[0] += 1
-        self._semaphore.acquire()
-        with self._lock:
-            self._counts[0] -= 1
-            self._counts[1] += 1
+        counts = self._counts
+        with self._cond:
+            counts[2] += 1
+            self._cond.wait_for(
+                lambda: counts[0] == 0 or counts[0] + nbytes <= self.budget_bytes
+            )
+            counts[2] -= 1
+            counts[0] += nbytes
+            counts[1] += 1
         try:
             yield
         finally:
-            with self._lock:
-                self._counts[1] -= 1
-            self._semaphore.release()
+            with self._cond:
+                counts[0] -= nbytes
+                counts[1] -= 1
+                self._cond.notify_all()
 
     def stats(self) -> dict[str, int]:
-        with self._lock:
+        with self._cond:
             return {
+                "budget_bytes": self.budget_bytes,
+                "reserved_bytes": self._counts[0],
+                "calls_active": self._counts[1],
+                "calls_waiting": self._counts[2],
                 "large_fit_slots": self.slots,
-                "large_fits_active": self._counts[1],
-                "large_fits_waiting": self._counts[0],
             }
diff --git a/tests/unit/test_host_memory.py b/tests/unit/test_host_memory.py
index cc4fd11..725a6c0 100644
--- a/tests/unit/test_host_memory.py
+++ b/tests/unit/test_host_memory.py
@@ -18,37 +18,73 @@ def test_slots_from_budget():
     assert HostMemoryGuard(10 * GiB, large_fit_bytes=48 * GiB).slots == 1  # never 0
 
 
-def test_small_fits_never_wait():
-    guard = HostMemoryGuard(GiB, large_fit_cells=100, large_fit_bytes=GiB)
-    with guard.reserve(False), guard.reserve(False):  # neither large: no slot used
-        assert guard.stats()["large_fits_active"] == 0
+def test_calls_estimated_at_zero_never_wait():
+    guard = HostMemoryGuard(GiB, large_fit_bytes=GiB)
+    with guard.reserve(0), guard.reserve(0):
+        assert guard.stats()["calls_active"] == 0
 
 
-def test_large_fits_are_limited_to_the_slots():
-    guard = HostMemoryGuard(
-        2 * GiB, large_fit_cells=100, large_fit_bytes=GiB
-    )  # 2 slots
-    inside, release = threading.Semaphore(0), threading.Event()
-
-    def large():
-        with guard.reserve(True):
+def _hold(guard, nbytes, inside, release):
+    def run():
+        with guard.reserve(nbytes):
             inside.release()
             release.wait(5)
 
-    threads = [threading.Thread(target=large, daemon=True) for _ in range(3)]
-    for t in threads:
-        t.start()
+    t = threading.Thread(target=run, daemon=True)
+    t.start()
+    return t
+
+
+def test_calls_are_admitted_while_their_bytes_fit_the_budget():
+    guard = HostMemoryGuard(10 * GiB, large_fit_bytes=GiB)
+    inside, release = threading.Semaphore(0), threading.Event()
+    threads = [_hold(guard, 4 * GiB, inside, release) for _ in range(3)]
     assert inside.acquire(timeout=5) and inside.acquire(timeout=5)
-    assert not inside.acquire(timeout=0.3)  # the third waits for a slot
-    assert guard.stats() == {
-        "large_fit_slots": 2,
-        "large_fits_active": 2,
-        "large_fits_waiting": 1,
-    }
+    assert not inside.acquire(timeout=0.3)  # 12 GiB > 10: the third waits
+    stats = guard.stats()
+    assert (stats["reserved_bytes"], stats["calls_active"], stats["calls_waiting"]) == (
+        8 * GiB,
+        2,
+        1,
+    )
     release.set()
     for t in threads:
         t.join(5)
-    assert guard.stats()["large_fits_active"] == 0
+    assert guard.stats()["reserved_bytes"] == 0
+
+
+def test_many_modest_calls_run_together():
+    """What the old slot guard got wrong: every TabPFN fit over 3M cells
+    took a whole 48 GiB slot though it uses a few GiB."""
+    guard = HostMemoryGuard(100 * GiB)
+    inside, release = threading.Semaphore(0), threading.Event()
+    threads = [_hold(guard, 5 * GiB, inside, release) for _ in range(8)]
+    assert all(inside.acquire(timeout=5) for _ in range(8))
+    release.set()
+    for t in threads:
+        t.join(5)
+
+
+def test_a_call_bigger_than_the_budget_runs_alone():
+    guard = HostMemoryGuard(10 * GiB, large_fit_bytes=GiB)
+    inside, release = threading.Semaphore(0), threading.Event()
+    big = _hold(guard, 50 * GiB, inside, release)
+    assert inside.acquire(timeout=5)  # admitted: nothing else was running
+    small = _hold(guard, GiB, inside, release)
+    assert not inside.acquire(timeout=0.3)  # waits for the big one
+    release.set()
+    for t in (big, small):
+        t.join(5)
+    assert inside.acquire(timeout=1)
+
+
+def test_a_multiprocessing_guard_reserves_and_reports():
+    import multiprocessing as mp
+
+    guard = HostMemoryGuard(10 * GiB, mp_context=mp.get_context("spawn"))
+    with guard.reserve(3 * GiB):
+        assert guard.stats()["reserved_bytes"] == 3 * GiB
+    assert guard.stats()["calls_active"] == 0
 
 
 class GatedFit(FakeBackend):
@@ -65,7 +101,7 @@ def fit(self, X, y, task):
 
 
 def test_a_waiting_large_fit_does_not_block_predicts_on_its_device():
-    guard = HostMemoryGuard(GiB, large_fit_cells=100, large_fit_bytes=GiB)  # 1 slot
+    guard = HostMemoryGuard(GiB, large_fit_cells=100, large_fit_bytes=GiB)  # 1 fit
     backends = [GatedFit(), GatedFit()]
     est = AdaptiveMemoryEstimator(
         fallback=PowerLawMemoryEstimator(A100_40GB_TABICL_CALIBRATION)
@@ -92,7 +128,7 @@ def test_a_waiting_large_fit_does_not_block_predicts_on_its_device():
         target=lambda: (engines[1].predict("small", [[1.0]]), done.set()), daemon=True
     ).start()
     assert done.wait(3), "a fit waiting on host memory blocked its device"
-    assert guard.stats()["large_fits_waiting"] == 1
+    assert guard.stats()["calls_waiting"] == 1
     backends[0].gate.set()
     backends[1].gate.set()
     t0.join(5)
@@ -178,11 +214,11 @@ def test_backend_hint_decides_and_predicts_are_guarded_too():
     engine = _engine(backend, guard)
     seen = []
     real = guard.reserve
-    guard.reserve = lambda large: (seen.append(large), real(large))[1]
+    guard.reserve = lambda nbytes: (seen.append(nbytes), real(nbytes))[1]
     engine.fit([[1.0, 2.0]] * 3, ["a"] * 3, dataset_id="d")
     engine.predict("d", [[1.0, 2.0]] * 5)
     # The hint, not the (huge) cell threshold, decided -- for fit AND predict.
-    assert seen == [True, True]
+    assert seen == [GiB, GiB]
     assert backend.queries == [(3, 2), (5, 2)]
 
 
@@ -191,7 +227,21 @@ def test_without_a_hint_the_cell_threshold_decides():
     engine = _engine(FakeBackend(bytes_hint=10), guard)
     seen = []
     real = guard.reserve
-    guard.reserve = lambda large: (seen.append(large), real(large))[1]
+    guard.reserve = lambda nbytes: (seen.append(nbytes), real(nbytes))[1]
     engine.fit([[1.0, 2.0]] * 3, ["a"] * 3, dataset_id="d")  # 6 cells: large
     engine.predict("d", [[1.0, 2.0]])  # 2 cells: not
-    assert seen == [True, False]
+    assert seen == [GiB, 0]
+
+
+def test_a_backend_estimate_wins_over_the_fallbacks():
+    class Estimating(_Offloading):
+        def host_bytes_estimate(self, n_rows, n_features):
+            return 7 * n_rows * n_features
+
+    guard = HostMemoryGuard(GiB, large_fit_cells=1, large_fit_bytes=GiB)
+    engine = _engine(Estimating(offloads=True), guard)
+    seen = []
+    real = guard.reserve
+    guard.reserve = lambda nbytes: (seen.append(nbytes), real(nbytes))[1]
+    engine.fit([[1.0, 2.0]] * 3, ["a"] * 3, dataset_id="d")
+    assert seen == [42]

From 613a2ac17fa19d02e3579f12623b97926b4af815 Mon Sep 17 00:00:00 2001
From: Amrit Krishnan <amrit110@gmail.com>
Date: Sat, 17 Oct 2026 22:59:15 -0400
Subject: [PATCH 26/39] fix(memory): FIFO node-wide host admission that
 survives worker death

Review of 0052e63..1131e41:
- A worker killed holding host bytes leaked them forever (every later
  large call on the node waited). Reservations are now per owner; the
  parent returns a dead worker's. The shared lock is taken with a
  timeout: a stuck lock fails calls and /readyz with a retryable 503.
- Admission is first come first served (no barging), so large fits
  aren't starved by modest ones; calls under 1 GiB never wait, so small
  TabPFN calls don't queue behind an over-budget call.
- At most `slots` offloading calls run at once, matching TabICL's
  cpu_safety_factor (0.6 / slots).
- At most one fit per device prepares ahead of the device, so queued
  fits don't each hold preprocessed copies in host memory.
- workers.py imports multiprocessing.util so its atexit order holds
  however it is imported; workers close together under one deadline.
- reserve() tolerates None and interrupted waits; estimates take the
  offload decision from the engine (one may_offload read per call).
- CHANGELOG: byte-based guard and the /readyz host_memory key rename.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 CHANGELOG.md                           |   4 +-
 src/tabctx/backends/tabicl.py          |  23 ++-
 src/tabctx/backends/tabpfn.py          |   4 +-
 src/tabctx/engine.py                   |  83 ++++++----
 src/tabctx/errors.py                   |   6 +
 src/tabctx/memory/host.py              | 184 +++++++++++++++++----
 src/tabctx/serve/app.py                |  13 +-
 src/tabctx/serve/factory.py            |   6 +-
 src/tabctx/serve/workers.py            |  57 +++++--
 tests/unit/test_engine_scheduling.py   |  22 +++
 tests/unit/test_host_memory.py         | 217 +++++++++++++++++++++++--
 tests/unit/test_tabicl_host_release.py |  13 ++
 tests/unit/test_workers.py             |  70 ++++++++
 13 files changed, 601 insertions(+), 101 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index b5b0d24..fc5552d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -14,7 +14,9 @@
 - ROADMAP: multi-GPU and CUDA-level track (#30-#37), PyTorch-level first.
 
 ### Changed
-- **TabICL fits prepare their CPU half before taking the device**: preprocessing (~30-37% of a fit) overlaps another fit's GPU half -- back-to-back fits 17-26% faster on an A100.
+- **Host memory is admitted by estimated bytes, node-wide** (`memory/host.py`). Each backend estimates a call's host use (TabICL ~5 KB per cell when it offloads stage outputs, else nothing; TabPFN ~0.6 KB per cell) and the guard admits calls against one budget shared by every worker process, first come first served (a large fit can't be starved by modest ones), with at most `slots` offloading calls at once (each budgets 1/slots of free host memory for staging). Calls under 1 GiB never wait. The old guard gave every call over 3M cells a whole 48 GiB slot, which serialized fits that use a few GiB (40k x 100 TabICL fits took p50 195 s under concurrent load on the 4x L4 rehearsal). A worker that dies holding host memory has it returned by its parent, and a stuck accounting lock fails calls with a retryable 503 instead of hanging. **`/readyz` `host_memory` keys changed**: `large_fits_active`/`large_fits_waiting` are now `calls_active`/`calls_waiting`, plus `reserved_bytes` and `calls_offloading`.
+- **TabICL fits prepare their CPU half before taking the device**: preprocessing (~30-37% of a fit) overlaps another fit's GPU half -- back-to-back fits 17-26% faster on an A100. At most one fit per device prepares ahead, so queued fits don't each hold preprocessed copies in host memory.
+- Worker processes shut down together under one 10 s deadline (was up to 10 s each, past a pod's grace period).
 - **tabicl's per-call log(n) host-to-device syncs are cached** (bit-identical): 10k x 50 fits 1.98 -> 1.66 s on an A100.
 - The cache's spill reads and writes run outside its lock, so bookkeeping and light predicts never wait on multi-GB disk I/O.
 - **TabPFN contexts share the model: ~42x less GPU memory per context.** TabPFN builds a fresh model for every estimator unless its built-model cache is on, so every cached context held its own ~850 MB copy of the weights on the GPU (measured on an A100: 4 contexts of a 300x6 table took 3.4 GB). `TabPFNBackend` now enables that cache (`TABPFN_MODEL_CACHE_SIZE`, default 16 entries, keyed by device and task): one shared model per (device, task), about 1.7 GB fixed per GPU, and a context costs only its own state (~20 MB for that table). Interleaved contexts predict bit-identically to each context alone.
diff --git a/src/tabctx/backends/tabicl.py b/src/tabctx/backends/tabicl.py
index a105a54..2f79877 100644
--- a/src/tabctx/backends/tabicl.py
+++ b/src/tabctx/backends/tabicl.py
@@ -72,8 +72,9 @@
 _restoring = threading.local()
 _warned = [False]  # one missing-hook warning per process
 # Offloaded stage-output bytes per cell (see TabICLBackend.may_offload).
-_HOST_BYTES_PER_OFFLOADED_CELL = 5_000
 _OUTPUT_BYTES_PER_CELL = 2560
+# Host bytes per cell of an offloading call (see host_bytes_estimate).
+_HOST_BYTES_PER_OFFLOADED_CELL = 5_000
 logger = logging.getLogger(__name__)
 _TASK_TAGS = {"classification": b"C", "regression": b"R"}
 
@@ -139,6 +140,9 @@ def install_cached_logn() -> bool:
     if original is None or getattr(original, "_tabctx_cached", False):
         return original is not None
 
+    # The cached tensors are created under whatever grad/inference mode the
+    # first call ran in; fine for serving (always inference), but anything
+    # that runs ssmax with autograd (fine-tuning) must not use this cache.
     @functools.lru_cache(maxsize=4096)
     def cached(n: int, device: Any, dtype: Any) -> Any:
         return original(n, device, dtype)
@@ -397,13 +401,16 @@ def may_offload(self, n_rows: int, n_features: int) -> bool:
         free, _ = torch.cuda.mem_get_info(torch.device(self._device))
         return _OUTPUT_BYTES_PER_CELL * n_rows * n_features > 0.4 * free
 
-    def host_bytes_estimate(self, n_rows: int, n_features: int) -> int:
-        """Host memory a call over n_rows x n_features may use: its offloaded
-        stage outputs when it may offload (measured ~46 GiB for 10M-cell
-        fits on an A100, ~4.9 KB per cell), else nothing worth reserving."""
-        if not self.may_offload(n_rows, n_features):
-            return 0
-        return _HOST_BYTES_PER_OFFLOADED_CELL * n_rows * n_features
+    def host_bytes_estimate(
+        self, n_rows: int, n_features: int, *, offloads: bool = False
+    ) -> int:
+        """Host memory a call over n_rows x n_features may use: when it
+        offloads, its staged stage outputs -- which scale with cells (the
+        offloaded tensors are per-cell activations); measured ~46 GiB for
+        10M-cell fits on an A100, ~4.9 KB per cell, budgeted at 5 KB. A
+        call that doesn't offload grows host RSS by < 0.5 GiB up to ~5M
+        cells: nothing worth reserving."""
+        return _HOST_BYTES_PER_OFFLOADED_CELL * n_rows * n_features if offloads else 0
 
     def _release_host_memory(self, model: Any) -> None:
         """Return the offload staging memory a large call leaves pinned.
diff --git a/src/tabctx/backends/tabpfn.py b/src/tabctx/backends/tabpfn.py
index f2d2e4a..624e4dd 100644
--- a/src/tabctx/backends/tabpfn.py
+++ b/src/tabctx/backends/tabpfn.py
@@ -162,7 +162,9 @@ def fit(self, X: ArrayLike, y: ArrayLike, task: Task) -> Any:
             self._shared_models.setdefault((task, index), module)
         return model
 
-    def host_bytes_estimate(self, n_rows: int, n_features: int) -> int:
+    def host_bytes_estimate(
+        self, n_rows: int, n_features: int, *, offloads: bool = False
+    ) -> int:
         """Host memory a fit or predict may use (preprocessing copies):
         measured on an A100 pod at most ~530 B per cell (100,000 x 100 grew
         RSS 4.9 GiB; 30M-cell tables 6-6.7 GiB), budgeted at 600."""
diff --git a/src/tabctx/engine.py b/src/tabctx/engine.py
index cbdb175..46a5949 100644
--- a/src/tabctx/engine.py
+++ b/src/tabctx/engine.py
@@ -36,6 +36,8 @@
 # A predict whose context is evicted again (under memory pressure) between
 # its restore and getting the device tries this many times.
 _PREDICT_ATTEMPTS = 3
+# Fits prepared (CPU half done) ahead of the one holding the device.
+_PREPARED_AHEAD = 1
 
 
 def _not_found(dataset_id: str) -> DatasetNotFoundError:
@@ -180,9 +182,11 @@ def __init__(
         # large predicts. _predict_lock runs one GPU predict at a time, so
         # a light predict (see LIGHT_PREDICT_ROWS) overlaps a running fit
         # rather than queueing behind it. The cache's own lock only guards
-        # bookkeeping. Lock order: host slot, _fit_lock, _predict_lock (the
-        # cache lock and _release_lock are leaves).
+        # bookkeeping. Lock order: prepare gate (fits only), host slot,
+        # _fit_lock, _predict_lock (the cache lock and _release_lock are
+        # leaves).
         self._fit_lock = threading.Lock()
+        self._prepare_gate = threading.BoundedSemaphore(1 + _PREPARED_AHEAD)
         # Re-entrant: a fit retrying alone holds it while the backend's
         # measurement takes it again (see _cuda.QuiescedMeasurement).
         self._predict_lock = threading.RLock()
@@ -317,33 +321,37 @@ def fit(
         # The fit's CPU half (backends that split it: TabICL's preprocessing,
         # ~1/3 of a fit) runs before taking the device, overlapping the GPU
         # half of whatever fit holds it now.
-        prepare = getattr(chosen, "prepare_fit", None)
-        prepared = prepare(X, y, task) if callable(prepare) else None
-        # Lock order everywhere: host slot, _fit_lock, _predict_lock. Host
-        # memory first, so a large fit waiting for a host slot doesn't hold
-        # this device meanwhile.
-        with self._host_slot(chosen, n_train, n_features), self._fit_lock:
-            self._make_room(estimator, estimated)
-            with self._measuring():
-                payload = self._fit_backend(
-                    chosen, X, y, task, estimator, estimated, prepared
+        # At most _PREPARED_AHEAD fits hold prepared (host-resident) data
+        # while waiting for the device, so queued fits don't each keep
+        # their preprocessed copies.
+        # Lock order everywhere: prepare gate, host slot, _fit_lock,
+        # _predict_lock. Host memory before the device, so a large fit
+        # waiting for host memory doesn't hold this device meanwhile.
+        with self._prepare_gate:
+            prepare = getattr(chosen, "prepare_fit", None)
+            prepared = prepare(X, y, task) if callable(prepare) else None
+            with self._host_slot(chosen, n_train, n_features), self._fit_lock:
+                self._make_room(estimator, estimated)
+                with self._measuring():
+                    payload = self._fit_backend(
+                        chosen, X, y, task, estimator, estimated, prepared
+                    )
+                    est_bytes = chosen.context_bytes_hint(n_train, n_features)
+                    if est_bytes is None:
+                        est_bytes = estimated
+                    else:
+                        self._learn(chosen, estimator, n_train, n_features, est_bytes)
+                context = CachedContext(
+                    dataset_id=resolved_id,
+                    backend_name=backend_name,
+                    task=task,
+                    n_train=n_train,
+                    n_features=n_features,
+                    payload=payload,
+                    est_bytes=est_bytes,
+                    feature_names=feature_names,
                 )
-                est_bytes = chosen.context_bytes_hint(n_train, n_features)
-                if est_bytes is None:
-                    est_bytes = estimated
-                else:
-                    self._learn(chosen, estimator, n_train, n_features, est_bytes)
-            context = CachedContext(
-                dataset_id=resolved_id,
-                backend_name=backend_name,
-                task=task,
-                n_train=n_train,
-                n_features=n_features,
-                payload=payload,
-                est_bytes=est_bytes,
-                feature_names=feature_names,
-            )
-            self._cache.put(context)
+                self._cache.put(context)
         return resolved_id
 
     def _check_feasible(
@@ -480,22 +488,25 @@ def _learn(
 
     def _host_slot(self, backend: TabularICLBackend, n_rows: int, n_features: int):
         """Reserve this call's estimated host memory against the process-
-        wide (or node-wide) budget: the backend's host_bytes_estimate, else
-        one large fit's worth when it may offload (may_offload, or the
-        guard's cell threshold). A no-op without a guard."""
+        wide (or node-wide) budget. Whether it offloads comes from the
+        backend's may_offload (else the guard's cell threshold); its bytes
+        from host_bytes_estimate, else one large fit's worth when it
+        offloads. A no-op without a guard."""
         guard = self._host_guard
         if guard is None:
             return contextlib.nullcontext()
-        estimate = getattr(backend, "host_bytes_estimate", None)
-        if callable(estimate):
-            return guard.reserve(estimate(n_rows, n_features))
         may_offload = getattr(backend, "may_offload", None)
-        large = (
+        offloads = bool(
             may_offload(n_rows, n_features)
             if callable(may_offload)
             else guard.is_large(n_rows, n_features)
         )
-        return guard.reserve(guard.large_fit_bytes if large else 0)
+        estimate = getattr(backend, "host_bytes_estimate", None)
+        if callable(estimate):
+            nbytes = estimate(n_rows, n_features, offloads=offloads)
+        else:
+            nbytes = guard.large_fit_bytes if offloads else 0
+        return guard.reserve(nbytes, offloads=offloads)
 
     def _restore(self, dataset_id: str) -> None:
         """Bring a spilled context back onto the device (DatasetNotFoundError
diff --git a/src/tabctx/errors.py b/src/tabctx/errors.py
index e0e36dc..646cd51 100644
--- a/src/tabctx/errors.py
+++ b/src/tabctx/errors.py
@@ -54,3 +54,9 @@ class BackendComputeError(TabctxError):
     """Wraps a backend-level compute failure (e.g. a CUDA OOM that slipped
     past admission control) so callers don't need to import torch to catch
     tabctx errors."""
+
+
+class HostMemoryUnavailableError(TabctxError):
+    """Raised when host-memory admission can't be decided right now (its
+    shared accounting is stuck, e.g. a worker died holding the lock);
+    retryable, so the API answers 503."""
diff --git a/src/tabctx/memory/host.py b/src/tabctx/memory/host.py
index 302eff7..e3be77b 100644
--- a/src/tabctx/memory/host.py
+++ b/src/tabctx/memory/host.py
@@ -13,14 +13,29 @@
 budget for the whole process -- or, built with mp_context, for every
 per-GPU worker process at once (host RAM is shared by the node). Each
 backend estimates its own calls (host_bytes_estimate: TabICL ~5 KB per
-cell when it may offload, TabPFN ~0.6 KB per cell); backends without an
-estimate count as one large fit above a cell threshold. A call waits
-until its bytes fit beside those already admitted (a call bigger than the
-whole budget runs alone, as the only one); calls estimated at 0 never
-wait. Calls that do offload budget their own staging against ``slots``,
-the number of large fits the budget holds (see TabICLBackend's
-cpu_safety_factor), so concurrent ones can't all claim the same free
-memory.
+cell when it offloads, TabPFN ~0.6 KB per cell); backends without an
+estimate count as one large fit above a cell threshold.
+
+Admission rules:
+- Calls under ``min_guarded_bytes`` (1 GiB, inside the reserve that
+  detect_host_budget_bytes already holds back) never wait, so small
+  predicts and fits stay lock-free however busy the guard is.
+- Others are admitted first come, first served: a call waits until it is
+  the oldest waiter AND its bytes fit beside those already admitted (a
+  call bigger than the whole budget runs alone). No barging, so a large
+  fit can't be starved by a stream of modest ones.
+- At most ``slots`` calls that offload run at once: the number of large
+  fits the budget holds. Each offloading call budgets its own staging
+  as a 1/slots share of free host memory (TabICLBackend's
+  cpu_safety_factor), so concurrent ones can't all claim the same memory.
+
+Robustness. Reservations are kept per owner (a worker process, or 0
+in-process), so when a worker dies holding bytes -- the host OOM kill
+this guard exists to prevent -- its parent returns them (release_owner)
+instead of leaking them until the replica restarts. The shared lock is
+taken with a timeout: a worker killed while holding it (a window of
+microseconds) makes calls fail with a retryable HostMemoryUnavailableError
+and /readyz report unhealthy, never hang.
 """
 
 from __future__ import annotations
@@ -32,12 +47,24 @@
 from pathlib import Path
 from typing import Any
 
+from tabctx.errors import HostMemoryUnavailableError
+
 GiB = 1024**3
 logger = logging.getLogger(__name__)
 # Calibrated on an A100-40GB (see module docstring): the fallback for
 # backends without a host_bytes_estimate.
 LARGE_FIT_CELLS = 3_000_000
 LARGE_FIT_HOST_BYTES = 48 * GiB
+# Calls below this never wait: detect_host_budget_bytes holds back a larger
+# reserve (8 GiB) for the rest of serving.
+MIN_GUARDED_BYTES = 1 * GiB
+LOCK_TIMEOUT_S = 10.0
+# Waiters re-check at least this often (a dead owner's release, a timeout).
+_POLL_S = 1.0
+# Per owner: reserved bytes, active calls, offloading calls, then queue
+# slots for waiting calls (a worker serves at most 16 calls at once).
+_QUEUE_SLOTS = 32
+_OWNER_BLOCK = 3 + _QUEUE_SLOTS
 
 
 CGROUP_LIMIT_PATHS = (
@@ -92,12 +119,24 @@ def _read_limit(limit_paths: tuple[str, ...], meminfo_path: str) -> int | None:
 
 
 class HostMemoryGuard:
+    """Node- or process-wide host memory admission (see module docstring).
+
+    Shared state (a plain list in-process, an mp Array across worker
+    processes): ``[next_ticket]`` then, per owner, ``[reserved bytes,
+    active calls, offloading calls, queue slots...]``. A waiting call
+    parks its ticket in one of its owner's queue slots; the oldest parked
+    ticket is the head of the queue.
+    """
+
     def __init__(
         self,
         budget_bytes: int,
         large_fit_cells: int = LARGE_FIT_CELLS,
         large_fit_bytes: int = LARGE_FIT_HOST_BYTES,
         mp_context: Any | None = None,
+        owners: int = 1,
+        min_guarded_bytes: int = MIN_GUARDED_BYTES,
+        lock_timeout_s: float = LOCK_TIMEOUT_S,
     ) -> None:
         self.budget_bytes = budget_bytes
         self.large_fit_bytes = large_fit_bytes
@@ -111,12 +150,24 @@ def __init__(
                 large_fit_bytes / GiB,
             )
         self._large_fit_cells = large_fit_cells
+        self._min_guarded = min_guarded_bytes
+        self._lock_timeout = lock_timeout_s
+        self._owners = max(1, owners)
+        self._owner = 0
+        size = 1 + self._owners * _OWNER_BLOCK
         if mp_context is None:
             self._cond: Any = threading.Condition()
-            self._counts: Any = [0, 0, 0]  # reserved bytes, active, waiting
+            self._state: Any = [0] * size
         else:
             self._cond = mp_context.Condition(mp_context.Lock())
-            self._counts = mp_context.Array("q", 3, lock=False)
+            self._state = mp_context.Array("q", size, lock=False)
+
+    def bind(self, owner: int) -> None:
+        """Account this process's calls to `owner` (a worker's index), so
+        its parent can release them if it dies (release_owner)."""
+        if not 0 <= owner < self._owners:
+            raise ValueError(f"owner {owner} out of range [0, {self._owners})")
+        self._owner = owner
 
     def is_large(self, n_rows: int, n_features: int) -> bool:
         """The fallback test for backends that can't estimate their host
@@ -124,34 +175,109 @@ def is_large(self, n_rows: int, n_features: int) -> bool:
         return n_rows * n_features >= self._large_fit_cells
 
     @contextlib.contextmanager
-    def reserve(self, nbytes: int) -> Iterator[None]:
-        """Hold `nbytes` of the host budget for the duration (no-op for 0)."""
-        if nbytes <= 0:
+    def reserve(self, nbytes: int | None, offloads: bool = False) -> Iterator[None]:
+        """Hold `nbytes` of the host budget (and, if the call offloads, one
+        of the `slots`) for the duration. Calls under min_guarded_bytes
+        that don't offload pass straight through."""
+        nbytes = max(0, int(nbytes or 0))
+        if nbytes < self._min_guarded and not offloads:
             yield
             return
-        counts = self._counts
-        with self._cond:
-            counts[2] += 1
-            self._cond.wait_for(
-                lambda: counts[0] == 0 or counts[0] + nbytes <= self.budget_bytes
-            )
-            counts[2] -= 1
-            counts[0] += nbytes
-            counts[1] += 1
+        state, base = self._state, self._block(self._owner)
+        with self._locked():
+            ticket = state[0] = state[0] + 1
+            slot = None
+            try:
+                while slot is None:
+                    slot = self._park(base, ticket)
+                    if slot is None:
+                        self._cond.wait(_POLL_S)
+                while not (self._head() == ticket and self._admits(nbytes, offloads)):
+                    self._cond.wait(_POLL_S)
+            finally:
+                if slot is not None:
+                    state[slot] = 0
+                self._cond.notify_all()  # the next waiter is now the head
+            state[base] += nbytes
+            state[base + 1] += 1
+            state[base + 2] += int(offloads)
         try:
             yield
         finally:
-            with self._cond:
-                counts[0] -= nbytes
-                counts[1] -= 1
+            with self._locked():
+                state[base] = max(0, state[base] - nbytes)
+                state[base + 1] = max(0, state[base + 1] - 1)
+                state[base + 2] = max(0, state[base + 2] - int(offloads))
                 self._cond.notify_all()
 
+    def release_owner(self, owner: int) -> None:
+        """Return everything a dead owner held or queued for."""
+        base = self._block(owner)
+        try:
+            with self._locked():
+                for i in range(base, base + _OWNER_BLOCK):
+                    self._state[i] = 0
+                self._cond.notify_all()
+        except HostMemoryUnavailableError:
+            logger.error("could not release host reservations of worker %d", owner)
+
     def stats(self) -> dict[str, int]:
-        with self._cond:
+        """Raises HostMemoryUnavailableError if the shared lock is stuck."""
+        with self._locked():
+            blocks = [self._block(o) for o in range(self._owners)]
             return {
                 "budget_bytes": self.budget_bytes,
-                "reserved_bytes": self._counts[0],
-                "calls_active": self._counts[1],
-                "calls_waiting": self._counts[2],
+                "reserved_bytes": sum(self._state[b] for b in blocks),
+                "calls_active": sum(self._state[b + 1] for b in blocks),
+                "calls_offloading": sum(self._state[b + 2] for b in blocks),
+                "calls_waiting": sum(
+                    1
+                    for b in blocks
+                    for i in range(b + 3, b + _OWNER_BLOCK)
+                    if self._state[i]
+                ),
                 "large_fit_slots": self.slots,
             }
+
+    # ---- internals (callers hold the lock) ----------------------------
+
+    @contextlib.contextmanager
+    def _locked(self) -> Iterator[None]:
+        if not self._cond.acquire(True, self._lock_timeout):
+            raise HostMemoryUnavailableError(
+                f"host memory accounting lock not acquired in {self._lock_timeout}s "
+                "(a worker may have died holding it); retry"
+            )
+        try:
+            yield
+        finally:
+            self._cond.release()
+
+    @staticmethod
+    def _block(owner: int) -> int:
+        return 1 + owner * _OWNER_BLOCK
+
+    def _park(self, base: int, ticket: int) -> int | None:
+        for i in range(base + 3, base + _OWNER_BLOCK):
+            if self._state[i] == 0:
+                self._state[i] = ticket
+                return i
+        return None
+
+    def _head(self) -> int:
+        return min(
+            (
+                self._state[i]
+                for o in range(self._owners)
+                for i in range(self._block(o) + 3, self._block(o) + _OWNER_BLOCK)
+                if self._state[i]
+            ),
+            default=0,
+        )
+
+    def _admits(self, nbytes: int, offloads: bool) -> bool:
+        blocks = [self._block(o) for o in range(self._owners)]
+        reserved = sum(self._state[b] for b in blocks)
+        if offloads and sum(self._state[b + 2] for b in blocks) >= self.slots:
+            return False
+        return reserved == 0 or reserved + nbytes <= self.budget_bytes
diff --git a/src/tabctx/serve/app.py b/src/tabctx/serve/app.py
index 4b59252..4529729 100644
--- a/src/tabctx/serve/app.py
+++ b/src/tabctx/serve/app.py
@@ -39,6 +39,7 @@
     BackendComputeError,
     CacheCapacityError,
     DatasetNotFoundError,
+    HostMemoryUnavailableError,
     InvalidInputError,
     UploadNotFoundError,
     UploadTooLargeError,
@@ -156,6 +157,7 @@ class TabctxPredictResponse(BaseModel):
 _ADMISSION_ERRORS = (AdmissionRejected, CacheCapacityError, UploadTooLargeError)
 _NOT_FOUND_ERRORS = (DatasetNotFoundError, UploadNotFoundError)
 _COMPUTE_ERRORS = (BackendComputeError,)
+_UNAVAILABLE_ERRORS = (HostMemoryUnavailableError,)
 
 
 def _map_error(e: Exception) -> HTTPException:
@@ -167,6 +169,8 @@ def _map_error(e: Exception) -> HTTPException:
         return HTTPException(413, str(e))
     if isinstance(e, _NOT_FOUND_ERRORS):
         return HTTPException(404, str(e))
+    if isinstance(e, _UNAVAILABLE_ERRORS):
+        return HTTPException(503, str(e))
     if isinstance(e, _COMPUTE_ERRORS):
         return HTTPException(507, str(e))
     raise e
@@ -284,6 +288,7 @@ def _legacy_predict_sync(self, req: LegacyPredictRequest) -> LegacyPredictRespon
             *_ADMISSION_ERRORS,
             *_NOT_FOUND_ERRORS,
             *_COMPUTE_ERRORS,
+            *_UNAVAILABLE_ERRORS,
         ) as e:
             logger.error("[%s] %s: %s", request_id, type(e).__name__, e)
             raise _map_error(e) from e
@@ -414,6 +419,7 @@ def _fit_sync(
             *_ADMISSION_ERRORS,
             *_NOT_FOUND_ERRORS,
             *_COMPUTE_ERRORS,
+            *_UNAVAILABLE_ERRORS,
         ) as e:
             raise _map_error(e) from e
         return FitResponse(
@@ -477,6 +483,7 @@ def _predict_sync(
             *_ADMISSION_ERRORS,
             *_NOT_FOUND_ERRORS,
             *_COMPUTE_ERRORS,
+            *_UNAVAILABLE_ERRORS,
         ) as e:
             raise _map_error(e) from e
         latency_ms = (time.monotonic() - start) * 1000
@@ -580,6 +587,10 @@ def readyz(self):
             reports = [w.report() for w in self._workers]
         except (WorkerDiedError, TimeoutError) as e:
             raise HTTPException(503, detail=str(e)) from e
+        try:
+            host_memory = self._host_guard.stats() if self._host_guard else None
+        except HostMemoryUnavailableError as e:
+            raise HTTPException(503, detail=str(e)) from e
         real_gpu_memory = None
         per_device: dict[str, dict] = {}
         if reports:
@@ -636,7 +647,7 @@ def readyz(self):
             "spill": self._spill_stats([r["spill"] for r in reports]),
             "real_gpu_memory": real_gpu_memory,
             "devices": self._device_breakdown(),
-            "host_memory": self._host_guard.stats() if self._host_guard else None,
+            "host_memory": host_memory,
         }
 
     def _cuda_devices(self) -> list[str]:
diff --git a/src/tabctx/serve/factory.py b/src/tabctx/serve/factory.py
index 4614894..8166cfa 100644
--- a/src/tabctx/serve/factory.py
+++ b/src/tabctx/serve/factory.py
@@ -713,7 +713,11 @@ def _build_process_engine(
     from tabctx.serve.workers import ProcessShard
 
     ctx = mp.get_context("spawn")
-    host_guard = HostMemoryGuard(host_budget, mp_context=ctx) if host_budget else None
+    host_guard = (
+        HostMemoryGuard(host_budget, mp_context=ctx, owners=len(devices))
+        if host_budget
+        else None
+    )
     shards = ProcessShard.start_all(
         settings,
         devices,
diff --git a/src/tabctx/serve/workers.py b/src/tabctx/serve/workers.py
index d13f335..a6278f7 100644
--- a/src/tabctx/serve/workers.py
+++ b/src/tabctx/serve/workers.py
@@ -40,6 +40,7 @@
 import atexit
 import itertools
 import logging
+import multiprocessing.util  # noqa: F401 -- registers its exit handler first
 import pickle
 import threading
 import time
@@ -74,15 +75,18 @@
 
 
 # Live shards, closed at interpreter exit -- before multiprocessing's own
-# exit handler (registered earlier, so run later) terminates the daemon
-# workers, which each reader thread would otherwise report as a death.
+# exit handler terminates the daemon workers, which each reader thread
+# would otherwise report as a death. atexit runs handlers in reverse, and
+# importing multiprocessing.util above registers mp's first, so ours runs
+# before it whatever imported this module.
 _LIVE_SHARDS: weakref.WeakSet = weakref.WeakSet()
+# How long a worker gets to exit after its pipe closes, before it is killed.
+_CLOSE_GRACE_S = 10.0
 
 
 @atexit.register
 def _close_live_shards() -> None:
-    for shard in list(_LIVE_SHARDS):
-        shard.close()
+    ProcessShard.close_all(list(_LIVE_SHARDS))
 
 
 class WorkerDiedError(BackendComputeError):
@@ -172,12 +176,15 @@ def _worker_main(
     n_devices: int,
     gpu_capacity: int | None,
     host_guard: Any,
+    index: int,
 ) -> None:
     """Entry point of a worker process: build the device's engine, then
     serve calls until the pipe closes."""
     logging.basicConfig(level=logging.INFO)
     from tabctx.serve.factory import build_shard
 
+    if host_guard is not None:
+        host_guard.bind(index)  # so the parent can release it if we die
     try:
         shard = build_shard(
             settings,
@@ -257,13 +264,16 @@ def __init__(
         gpu_capacity: int | None,
         host_guard: Any,
         mp_context: Any,
+        index: int = 0,
     ) -> None:
         self.device = device
+        self._index = index
+        self._host_guard = host_guard
         parent, child = mp_context.Pipe()
         self._conn = parent
         self._proc = mp_context.Process(
             target=_worker_main,
-            args=(child, settings, device, n_devices, gpu_capacity, host_guard),
+            args=(child, settings, device, n_devices, gpu_capacity, host_guard, index),
             name=f"tabctx-worker-{device}",
             daemon=True,
         )
@@ -327,20 +337,35 @@ def start_all(
                 gpu_capacity=gpu_capacity,
                 host_guard=host_guard,
                 mp_context=mp_context,
+                index=i,
             )
-            for d in devices
+            for i, d in enumerate(devices)
         ]
         try:
             deadline = time.monotonic() + ready_timeout_s
             for s in shards:
                 s.wait_ready(max(1.0, deadline - time.monotonic()))
         except BaseException:
-            for s in shards:
-                s.close()
+            cls.close_all(shards)
             raise
         return shards
 
+    @staticmethod
+    def close_all(shards: list[ProcessShard]) -> None:
+        """Shut workers down together: ask them all to exit, then wait with
+        one shared deadline (closing them one by one could take 10 s each,
+        past a pod's termination grace period)."""
+        for shard in shards:
+            shard._request_exit()
+        deadline = time.monotonic() + _CLOSE_GRACE_S
+        for shard in shards:
+            shard._finish_close(max(0.0, deadline - time.monotonic()))
+
     def close(self) -> None:
+        self._request_exit()
+        self._finish_close(_CLOSE_GRACE_S)
+
+    def _request_exit(self) -> None:
         with self._state_lock:
             if self._dead is None:
                 self._dead = f"worker for {self.device} was shut down"
@@ -349,13 +374,21 @@ def close(self) -> None:
                 self._conn.send(None)
             except (OSError, ValueError):
                 pass
-        self._proc.join(10)
+
+    def _finish_close(self, grace_s: float) -> None:
+        self._proc.join(grace_s)
         if self._proc.is_alive():
             self._proc.kill()
             self._proc.join(5)
         self._conn.close()
+        self._release_host_memory()
         self._fail_pending(f"worker for {self.device} was shut down")
 
+    def _release_host_memory(self) -> None:
+        """Return host memory the (now exited) worker had reserved."""
+        if self._host_guard is not None and not self._proc.is_alive():
+            self._host_guard.release_owner(self._index)
+
     @property
     def alive(self) -> bool:
         return self._dead is None and self._proc.is_alive()
@@ -390,7 +423,11 @@ def _read_loop(self) -> None:
                 fut.set_result(payload)
             else:
                 fut.set_exception(payload)
-        self._proc.join(1)  # reap it, for a meaningful exit code
+        self._proc.join(5)  # reap it, for a meaningful exit code
+        if self._dead is None and self._proc.is_alive():
+            self._proc.kill()  # its pipe is gone: it can't serve, and must
+            self._proc.join(5)  # exit before its reservations are returned
+        self._release_host_memory()
         reason = (
             f"worker for {self.device} exited (code {self._proc.exitcode}); "
             "its contexts are lost"
diff --git a/tests/unit/test_engine_scheduling.py b/tests/unit/test_engine_scheduling.py
index 7a0d179..a1803aa 100644
--- a/tests/unit/test_engine_scheduling.py
+++ b/tests/unit/test_engine_scheduling.py
@@ -349,3 +349,25 @@ def test_an_infeasible_fit_is_rejected_before_preparing():
     with pytest.raises(AdmissionRejected):
         engine.fit(*table(2000), dataset_id="huge")  # 20,000 > 10,000 bytes
     assert backend.prepared_during == []
+
+
+def test_only_one_fit_prepares_ahead_of_the_device():
+    """Prepared data lives in host memory until its fit gets the device:
+    queued fits beyond one ahead wait before preparing, so N queued large
+    tables don't hold N preprocessed copies at once."""
+    backend = PreparingFake()
+    engine = TabctxEngine(
+        backend=backend, estimator=CellEstimator(), cache=ContextCacheManager(CAPACITY)
+    )
+    backend.engine = engine
+    backend.fit_gate = threading.Event()
+    threads = [start(engine.fit, *table(10), dataset_id=f"d{i}")[0] for i in range(4)]
+    assert backend.fit_entered.wait(5)
+    for _ in range(50):  # give queued fits every chance to prepare
+        threading.Event().wait(0.01)
+    assert len(backend.prepared_during) == 2  # the running fit + one ahead
+    backend.fit_gate.set()
+    for t in threads:
+        t.join(5)
+    assert len(backend.prepared_during) == 4
+    assert all(engine._cache.holds(f"d{i}") for i in range(4))
diff --git a/tests/unit/test_host_memory.py b/tests/unit/test_host_memory.py
index 725a6c0..9c02b2f 100644
--- a/tests/unit/test_host_memory.py
+++ b/tests/unit/test_host_memory.py
@@ -1,6 +1,9 @@
 """Unit tests for process-wide host-memory admission (memory/host.py)."""
 
 import threading
+import time
+
+import pytest
 
 from tabctx.backends.fake import FakeBackend
 from tabctx.cache.manager import ContextCacheManager
@@ -87,6 +90,188 @@ def test_a_multiprocessing_guard_reserves_and_reports():
     assert guard.stats()["calls_active"] == 0
 
 
+def test_calls_under_the_floor_never_wait_even_beside_a_call_running_alone():
+    """An over-budget call runs alone, but small calls (inside the 8 GiB
+    reserve) must keep flowing: they bypass the guard entirely."""
+    guard = HostMemoryGuard(10 * GiB, large_fit_bytes=GiB)
+    inside, release = threading.Semaphore(0), threading.Event()
+    big = _hold(guard, 50 * GiB, inside, release)
+    assert inside.acquire(timeout=5)
+    with guard.reserve(GiB // 2):  # immediate
+        assert guard.stats()["calls_active"] == 1  # not even counted
+    release.set()
+    big.join(5)
+
+
+def test_admission_is_first_come_first_served():
+    """A large call queued behind running ones isn't overtaken by smaller
+    ones that would fit now: without this, steady modest traffic starves
+    it forever."""
+    guard = HostMemoryGuard(10 * GiB, large_fit_bytes=GiB)
+    inside, release = threading.Semaphore(0), threading.Event()
+    running = _hold(guard, 4 * GiB, inside, release)
+    assert inside.acquire(timeout=5)
+    order: list[str] = []
+    big_release = threading.Event()
+
+    def call(name, nbytes, done):
+        with guard.reserve(nbytes):
+            order.append(name)
+            done.wait(5)
+
+    big = threading.Thread(target=call, args=("big", 8 * GiB, big_release), daemon=True)
+    big.start()
+    _wait_for(lambda: guard.stats()["calls_waiting"] == 1)
+    small = threading.Thread(target=call, args=("small", 3 * GiB, release), daemon=True)
+    small.start()  # 4 + 3 <= 10 would fit, but "big" was first
+    _wait_for(lambda: guard.stats()["calls_waiting"] == 2)
+    assert order == []
+    release.set()  # the running call ends: big goes first (8 + 3 > 10)
+    _wait_for(lambda: order == ["big"])
+    assert not _settles(lambda: len(order) > 1)
+    big_release.set()
+    for t in (running, big, small):
+        t.join(5)
+    assert order == ["big", "small"]
+
+
+def test_offloading_calls_are_capped_at_slots():
+    """Each offloading call budgets 1/slots of free host memory for its
+    staging, so no more than `slots` may run at once even if bytes fit."""
+    guard = HostMemoryGuard(100 * GiB, large_fit_bytes=50 * GiB)  # 2 slots
+    held = threading.Event()
+    entered = threading.Semaphore(0)
+
+    def offloading():
+        with guard.reserve(10 * GiB, offloads=True):
+            entered.release()
+            held.wait(5)
+
+    threads = [threading.Thread(target=offloading, daemon=True) for _ in range(3)]
+    for t in threads:
+        t.start()
+    assert entered.acquire(timeout=5) and entered.acquire(timeout=5)
+    assert not entered.acquire(timeout=0.3)
+    assert guard.stats()["calls_offloading"] == 2
+    held.set()
+    assert entered.acquire(timeout=5)
+    for t in threads:
+        t.join(5)
+
+
+def test_a_none_estimate_counts_as_nothing():
+    guard = HostMemoryGuard(GiB)
+    with guard.reserve(None):
+        assert guard.stats()["calls_active"] == 0
+
+
+def test_an_interrupted_waiter_leaves_the_queue():
+    guard = HostMemoryGuard(10 * GiB, large_fit_bytes=GiB)
+    inside, release = threading.Semaphore(0), threading.Event()
+    running = _hold(guard, 8 * GiB, inside, release)
+    assert inside.acquire(timeout=5)
+    real_wait = guard._cond.wait
+
+    def interrupted(timeout=None):
+        guard._cond.wait = real_wait
+        raise KeyboardInterrupt
+
+    guard._cond.wait = interrupted
+    try:
+        with guard.reserve(4 * GiB):
+            raise AssertionError("must not be admitted")
+    except KeyboardInterrupt:
+        pass
+    stats = guard.stats()
+    assert (stats["calls_waiting"], stats["calls_active"]) == (0, 1)
+    release.set()
+    running.join(5)
+
+
+def test_a_stuck_lock_fails_retryably_instead_of_hanging():
+    from tabctx.errors import HostMemoryUnavailableError
+
+    guard = HostMemoryGuard(10 * GiB, large_fit_bytes=GiB, lock_timeout_s=0.2)
+    holding, done = threading.Event(), threading.Event()
+
+    def hold_lock():  # stands in for a worker killed holding it
+        with guard._cond:
+            holding.set()
+            done.wait(5)
+
+    t = threading.Thread(target=hold_lock, daemon=True)
+    t.start()
+    assert holding.wait(5)
+    with pytest.raises(HostMemoryUnavailableError), guard.reserve(2 * GiB):
+        pass
+    with pytest.raises(HostMemoryUnavailableError):
+        guard.stats()
+    done.set()
+    t.join(5)
+
+
+def _reserve_in_child(guard, owner, nbytes, ready):
+    guard.bind(owner)
+    with guard.reserve(nbytes):
+        ready.set()
+        threading.Event().wait(60)  # until killed
+
+
+def test_reservations_are_node_wide_and_a_dead_owner_is_released():
+    import multiprocessing as mp
+
+    ctx = mp.get_context("spawn")
+    guard = HostMemoryGuard(10 * GiB, large_fit_bytes=GiB, mp_context=ctx, owners=2)
+    ready = ctx.Event()
+    child = ctx.Process(
+        target=_reserve_in_child, args=(guard, 1, 8 * GiB, ready), daemon=True
+    )
+    child.start()
+    try:
+        assert ready.wait(60)
+        assert guard.stats()["reserved_bytes"] == 8 * GiB  # seen by the parent
+        admitted = threading.Event()
+
+        def parent_call():
+            with guard.reserve(4 * GiB):
+                admitted.set()
+
+        t = threading.Thread(target=parent_call, daemon=True)
+        t.start()
+        assert not admitted.wait(0.5)  # the child's bytes block it
+        child.kill()  # OOM-killed: its finally never runs
+        child.join(10)
+        assert not admitted.wait(0.5)  # leaked until released...
+        guard.release_owner(1)
+        assert admitted.wait(5)  # ...then returned
+        t.join(5)
+        assert guard.stats()["reserved_bytes"] == 0
+    finally:
+        if child.is_alive():
+            child.kill()
+
+
+def test_bind_rejects_unknown_owners():
+    with pytest.raises(ValueError):
+        HostMemoryGuard(GiB, owners=2).bind(2)
+
+
+def _settles(predicate, within=0.3):
+    deadline = time.monotonic() + within
+    while time.monotonic() < deadline:
+        if predicate():
+            return True
+        time.sleep(0.01)
+    return False
+
+
+def _wait_for(predicate, timeout=5.0):
+    deadline = time.monotonic() + timeout
+    while not predicate():
+        assert time.monotonic() < deadline, "condition not reached"
+        time.sleep(0.01)
+
+
 class GatedFit(FakeBackend):
     def __init__(self) -> None:
         super().__init__(bytes_hint=10)
@@ -196,6 +381,16 @@ def may_offload(self, n_rows, n_features):
         return self.offloads
 
 
+def _spy(guard):
+    seen = []
+    real = guard.reserve
+    guard.reserve = lambda nbytes, offloads=False: (
+        seen.append((nbytes, offloads)),
+        real(nbytes, offloads=offloads),
+    )[1]
+    return seen
+
+
 def _engine(backend, guard):
     est = AdaptiveMemoryEstimator(
         fallback=PowerLawMemoryEstimator(A100_40GB_TABICL_CALIBRATION)
@@ -212,36 +407,30 @@ def test_backend_hint_decides_and_predicts_are_guarded_too():
     guard = HostMemoryGuard(GiB, large_fit_cells=10**12, large_fit_bytes=GiB)
     backend = _Offloading(offloads=True)
     engine = _engine(backend, guard)
-    seen = []
-    real = guard.reserve
-    guard.reserve = lambda nbytes: (seen.append(nbytes), real(nbytes))[1]
+    seen = _spy(guard)
     engine.fit([[1.0, 2.0]] * 3, ["a"] * 3, dataset_id="d")
     engine.predict("d", [[1.0, 2.0]] * 5)
     # The hint, not the (huge) cell threshold, decided -- for fit AND predict.
-    assert seen == [GiB, GiB]
+    assert seen == [(GiB, True), (GiB, True)]
     assert backend.queries == [(3, 2), (5, 2)]
 
 
 def test_without_a_hint_the_cell_threshold_decides():
     guard = HostMemoryGuard(GiB, large_fit_cells=6, large_fit_bytes=GiB)
     engine = _engine(FakeBackend(bytes_hint=10), guard)
-    seen = []
-    real = guard.reserve
-    guard.reserve = lambda nbytes: (seen.append(nbytes), real(nbytes))[1]
+    seen = _spy(guard)
     engine.fit([[1.0, 2.0]] * 3, ["a"] * 3, dataset_id="d")  # 6 cells: large
     engine.predict("d", [[1.0, 2.0]])  # 2 cells: not
-    assert seen == [GiB, 0]
+    assert seen == [(GiB, True), (0, False)]
 
 
 def test_a_backend_estimate_wins_over_the_fallbacks():
     class Estimating(_Offloading):
-        def host_bytes_estimate(self, n_rows, n_features):
-            return 7 * n_rows * n_features
+        def host_bytes_estimate(self, n_rows, n_features, *, offloads=False):
+            return 7 * n_rows * n_features * (2 if offloads else 1)
 
     guard = HostMemoryGuard(GiB, large_fit_cells=1, large_fit_bytes=GiB)
     engine = _engine(Estimating(offloads=True), guard)
-    seen = []
-    real = guard.reserve
-    guard.reserve = lambda nbytes: (seen.append(nbytes), real(nbytes))[1]
+    seen = _spy(guard)
     engine.fit([[1.0, 2.0]] * 3, ["a"] * 3, dataset_id="d")
-    assert seen == [42]
+    assert seen == [(84, True)]
diff --git a/tests/unit/test_tabicl_host_release.py b/tests/unit/test_tabicl_host_release.py
index 09b5445..f27b48f 100644
--- a/tests/unit/test_tabicl_host_release.py
+++ b/tests/unit/test_tabicl_host_release.py
@@ -75,3 +75,16 @@ def test_pools_are_cleared_even_when_the_fit_fails():
     with pytest.raises(RuntimeError):
         _backend(est).fit([[1.0]], [0], "classification")
     assert pools(est) == [1, 1]
+
+
+def test_host_estimate_reserves_only_for_offloading_calls():
+    b = object.__new__(TabICLBackend)
+    assert b.host_bytes_estimate(100_000, 100, offloads=False) == 0
+    assert b.host_bytes_estimate(100_000, 100, offloads=True) == 5_000 * 10_000_000
+
+
+def test_tabpfn_host_estimate_scales_with_cells():
+    from tabctx.backends.tabpfn import TabPFNBackend
+
+    b = object.__new__(TabPFNBackend)
+    assert b.host_bytes_estimate(100_000, 100) == 600 * 10_000_000
diff --git a/tests/unit/test_workers.py b/tests/unit/test_workers.py
index f76488a..68fbb9d 100644
--- a/tests/unit/test_workers.py
+++ b/tests/unit/test_workers.py
@@ -138,3 +138,73 @@ def __init__(self, a, b):
     portable = _portable(TwoArgs(1, 2))
     assert type(portable) is RuntimeError and "TwoArgs" in str(portable)
     assert isinstance(_portable(DatasetNotFoundError("x")), DatasetNotFoundError)
+
+
+def test_a_dead_workers_host_reservations_are_returned():
+    """A worker OOM-killed mid-fit never runs its release: its parent must
+    return the bytes, or every later large call on the node waits forever."""
+    import multiprocessing as mp
+
+    from tabctx.memory.host import GiB, HostMemoryGuard
+
+    ctx = mp.get_context("spawn")
+    guard = HostMemoryGuard(10 * GiB, mp_context=ctx, owners=2)
+    shards = ProcessShard.start_all(
+        _settings(),
+        ("cuda:0", "cuda:1"),
+        gpu_capacity=None,
+        host_guard=guard,
+        mp_context=ctx,
+    )
+    try:
+        victim = shards[1]
+        guard._state[guard._block(1)] = 8 * GiB  # what it held when killed
+        guard._state[guard._block(1) + 1] = 1
+        victim._proc.kill()
+        for _ in range(500):
+            if guard.stats()["reserved_bytes"] == 0:
+                break
+            threading.Event().wait(0.01)
+        assert guard.stats()["reserved_bytes"] == 0
+        assert guard.stats()["calls_active"] == 0
+    finally:
+        ProcessShard.close_all(shards)
+
+
+def test_workers_close_together_within_one_grace_period(monkeypatch):
+    import time
+
+    import tabctx.serve.workers as workers
+
+    b = build_engine(
+        ServeSettings(backends=("fake",), devices=("cuda:0", "cuda:1", "cuda:2"))
+    )
+    shards = list(b.worker_shards)
+    monkeypatch.setattr(workers, "_CLOSE_GRACE_S", 1.0)
+    for s in shards:  # workers that ignore the exit request
+        monkeypatch.setattr(s, "_request_exit", lambda: None)
+    start = time.monotonic()
+    ProcessShard.close_all(shards)
+    # One shared deadline (~1 s), not 1 s per worker, then killed.
+    assert time.monotonic() - start < 2.5
+    assert not any(s._proc.is_alive() for s in shards)
+
+
+def test_mp_exit_handler_is_registered_before_ours():
+    """Our atexit handler must run first (atexit is LIFO), however the
+    module was imported: workers.py imports multiprocessing.util itself."""
+    import subprocess
+    import sys
+
+    out = subprocess.run(
+        [
+            sys.executable,
+            "-c",
+            "import sys, tabctx.serve.workers; "
+            "print('multiprocessing.util' in sys.modules)",
+        ],
+        capture_output=True,
+        text=True,
+        check=True,
+    )
+    assert out.stdout.strip() == "True"

From ae441e8f528598f399d392d4000710e9f4b65db3 Mon Sep 17 00:00:00 2001
From: Amrit Krishnan <amrit110@gmail.com>
Date: Sat, 17 Oct 2026 23:28:34 -0400
Subject: [PATCH 27/39] fix(serve): export HostMemoryUnavailableError (the app
 failed to start)

serve/app.py imports it from the package root; unit tests skip the app
without Ray, so only the replica start caught it. test_serve_imports
checks every tabctx name the serve layer imports, without Ray.

Also benchmarks/bench_compile.py (torch.compile of TabICL's attention
blocks: GPU half, peak, accuracy vs eager).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 benchmarks/bench_compile.py      | 154 +++++++++++++++++++++++++++++++
 src/tabctx/__init__.py           |   2 +
 tests/unit/test_serve_imports.py |  39 ++++++++
 3 files changed, 195 insertions(+)
 create mode 100644 benchmarks/bench_compile.py
 create mode 100644 tests/unit/test_serve_imports.py

diff --git a/benchmarks/bench_compile.py b/benchmarks/bench_compile.py
new file mode 100644
index 0000000..b078658
--- /dev/null
+++ b/benchmarks/bench_compile.py
@@ -0,0 +1,154 @@
+#!/usr/bin/env python3
+"""Does torch.compile speed up TabICL's GPU half, and is it exact enough? (#32/#36)
+
+The fit profile (#31) put layer norm at ~21% of the GPU half -- memory-bound
+passes over [ensembles, features, rows, 128] activations, each next to a
+residual add or a GELU MLP that also streams the tensor. Inductor can fuse
+those elementwise passes. This compares, on one device:
+
+- eager (baseline);
+- ``block``: each attention block's whole forward compiled (norms,
+  residual adds and the GELU MLP fused around the attention kernel).
+
+For each it reports the GPU half's time (TabICLBackend.fit_prepared, CPU
+preprocessing excluded), the first call's compile cost, predict latency, and
+the largest probability difference against eager on held-out rows.
+
+Usage (on a GPU host with tabicl and the weights):
+  bench_compile.py --device cuda:0 --shapes 10000x50 40000x100 --repeats 3 \
+      --variant block --baseline eager.npz --out compile-a100.jsonl
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+import time
+
+import numpy as np
+import torch
+
+
+def parse_shape(text: str) -> tuple[int, int]:
+    rows, features = text.lower().split("x")
+    return int(rows), int(features)
+
+
+def make_table(rows: int, features: int, seed: int):
+    """A noisy nonlinear task (so accuracy is informative, not saturated)."""
+    rng = np.random.default_rng(seed)
+    X = rng.normal(size=(rows, features)).astype(np.float32)
+    j = min(2, features - 1)
+    signal = X[:, 0] + 0.5 * X[:, min(1, features - 1)] * X[:, j] + np.sin(2 * X[:, j])
+    y = (signal + rng.normal(scale=0.8, size=rows) > 0).astype(int)
+    return X, y
+
+
+def install(variant: str, mode: str) -> None:
+    """Patch tabicl's attention block class (all instances share one graph:
+    nn-module parameters are graph inputs, so instances don't recompile)."""
+    from tabicl._model.layers import MultiheadAttentionBlock as Block
+
+    if variant == "eager":
+        return
+    if variant == "block":
+        # Blocks see several static shapes (column vs row vs ICL stages,
+        # induced points vs features): allow a graph per shape family.
+        torch._dynamo.config.recompile_limit = 64
+        Block.forward = torch.compile(Block.forward, dynamic=True, mode=mode)
+        return
+    raise SystemExit(f"unknown variant {variant}")
+
+
+def run(device: str, rows: int, features: int, repeats: int, amp: str) -> dict:
+    from tabctx.backends.tabicl import TabICLBackend
+
+    backend = TabICLBackend(device=device)
+    if amp != "auto":  # tabicl's default: fp16 autocast above 1024 rows
+        make = backend._make_estimator
+
+        def make_with_amp(task):
+            est = make(task)
+            est.use_amp = amp == "on"
+            return est
+
+        backend._make_estimator = make_with_amp
+    backend.warmup()
+    torch.cuda.reset_peak_memory_stats(device)
+    X, y = make_table(rows, features, seed=0)
+    X_test, y_test = make_table(2048, features, seed=1)
+    times, payload = [], None
+    first = 0.0
+    for i in range(repeats + 1):
+        prepared = backend.prepare_fit(X, y, "classification")
+        torch.cuda.synchronize(device)
+        t0 = time.perf_counter()
+        payload = backend.fit_prepared(prepared)
+        torch.cuda.synchronize(device)
+        elapsed = time.perf_counter() - t0
+        if i == 0:
+            first = elapsed
+        else:
+            times.append(elapsed)
+    backend.predict(payload, X_test, return_proba=True)  # warm the predict path
+    torch.cuda.synchronize(device)
+    t0 = time.perf_counter()
+    outcome = backend.predict(payload, X_test, return_proba=True)
+    torch.cuda.synchronize(device)
+    proba = np.asarray(outcome.probabilities, dtype=np.float64)
+    p_true = np.clip(proba[np.arange(len(y_test)), y_test], 1e-12, 1.0)
+    return {
+        "accuracy": round(float((proba.argmax(1) == y_test).mean()), 4),
+        "log_loss": round(float(-np.log(p_true).mean()), 5),
+        "first_s": round(first, 3),
+        "gpu_half_s": round(float(np.median(times)), 3),
+        "predict_s": round(time.perf_counter() - t0, 3),
+        "proba": proba,
+        "peak_gib": round(torch.cuda.max_memory_allocated(device) / 1024**3, 2),
+    }
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--device", default="cuda:0")
+    ap.add_argument("--shapes", nargs="+", default=["10000x50", "40000x100"])
+    ap.add_argument("--repeats", type=int, default=3)
+    ap.add_argument("--variant", default="eager")
+    ap.add_argument("--mode", default="default")
+    ap.add_argument("--amp", choices=("auto", "on", "off"), default="auto")
+    ap.add_argument("--baseline", help="npz of eager probabilities to compare against")
+    ap.add_argument("--save-proba", help="write this run's probabilities (npz)")
+    ap.add_argument("--out")
+    args = ap.parse_args()
+
+    install(args.variant, args.mode)
+    baseline = dict(np.load(args.baseline)) if args.baseline else {}
+    saved = {}
+    for text in args.shapes:
+        rows, features = parse_shape(text)
+        result = run(args.device, rows, features, args.repeats, args.amp)
+        proba = result.pop("proba")
+        saved[text] = proba
+        if text in baseline:
+            ref = baseline[text]
+            result["max_abs_proba_diff"] = float(np.abs(proba - ref).max())
+            result["label_agreement"] = float((proba.argmax(1) == ref.argmax(1)).mean())
+        record = {
+            "variant": args.variant,
+            "mode": args.mode,
+            "amp": args.amp,
+            "shape": text,
+            **result,
+        }
+        print(json.dumps(record), flush=True)
+        if args.out:
+            with open(args.out, "a") as f:
+                f.write(json.dumps(record) + "\n")
+    if args.save_proba:
+        np.savez(args.save_proba, **saved)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/src/tabctx/__init__.py b/src/tabctx/__init__.py
index 45b0846..4827664 100644
--- a/src/tabctx/__init__.py
+++ b/src/tabctx/__init__.py
@@ -16,6 +16,7 @@
     BackendComputeError,
     CacheCapacityError,
     DatasetNotFoundError,
+    HostMemoryUnavailableError,
     InvalidInputError,
     TabctxError,
     UploadNotFoundError,
@@ -31,6 +32,7 @@
     "ContextCacheManager",
     "DatasetNotFoundError",
     "EngineStats",
+    "HostMemoryUnavailableError",
     "InvalidInputError",
     "PredictOutcome",
     "TabctxEngine",
diff --git a/tests/unit/test_serve_imports.py b/tests/unit/test_serve_imports.py
new file mode 100644
index 0000000..a92af82
--- /dev/null
+++ b/tests/unit/test_serve_imports.py
@@ -0,0 +1,39 @@
+"""The serving app imports names from tabctx modules; check they exist
+without importing the app itself (which needs Ray, absent from the unit
+environment -- a missing export once passed every unit test and only
+failed when Ray Serve started the replica)."""
+
+import ast
+import importlib
+from pathlib import Path
+
+import pytest
+
+SERVE = Path(__file__).resolve().parents[2] / "src" / "tabctx" / "serve"
+
+
+def _tabctx_imports(path: Path):
+    for node in ast.walk(ast.parse(path.read_text())):
+        if isinstance(node, ast.ImportFrom) and (node.module or "").startswith(
+            "tabctx"
+        ):
+            for alias in node.names:
+                yield node.module, alias.name
+
+
+@pytest.mark.parametrize("path", sorted(SERVE.glob("*.py")), ids=lambda p: p.name)
+def test_every_tabctx_name_the_serve_layer_imports_exists(path):
+    missing = []
+    for module, name in _tabctx_imports(path):
+        try:
+            mod = importlib.import_module(module)
+        except ImportError as e:  # an optional dependency of that module
+            if "tabctx" in str(e):
+                missing.append(f"{module} ({e})")
+            continue
+        if not hasattr(mod, name):
+            try:
+                importlib.import_module(f"{module}.{name}")
+            except ImportError:
+                missing.append(f"{module}.{name}")
+    assert not missing, missing

From e7ae9c1b1f807ac0503924bbdb50d9bbfb3e80f6 Mon Sep 17 00:00:00 2001
From: Amrit Krishnan <amrit110@gmail.com>
Date: Sat, 17 Oct 2026 23:52:46 -0400
Subject: [PATCH 28/39] perf(tabicl): torch.compile the attention blocks
 (1.7-2.2x GPU half)

Inductor fuses the memory-bound layer norms, residual adds and GELU MLP
around each attention block (layer norm was ~21% of the GPU half, #31).
A100, tabicl defaults (fp16 autocast), benchmarks/bench_compile.py:
  10000x50   0.91 -> 0.52 s, peak 9.9 -> 7.3 GiB
  40000x100  7.65 -> 4.30 s, peak 33.1 -> 27.0 GiB
  20000x200  6.51 -> 2.92 s, peak 31.1 -> 23.4 GiB
Accuracy/AUC/R^2 unchanged on the evaluation datasets (eval_kv_precision
--compile), 0 wrong answers with predicts beside a fit (bench_overlap
--compile), 100-row predicts 20-35% faster.

Dynamic shapes needed tabicl's ssmax log(n) computed symbolically inside
compiled code: math.log on a Python int specialized every graph to one
row count (each new table recompiled for ~11 s). Graphs per call family
(~40) compile in warmup (~75 s per worker, in parallel). Compile errors
fall back to eager. TABCTX_COMPILE=0 turns it off.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 benchmarks/bench_compile.py                  | 19 ++++++
 benchmarks/bench_overlap.py                  |  5 +-
 benchmarks/eval_kv_precision.py              | 12 +++-
 benchmarks/profile_fit_predict.py            |  6 +-
 src/tabctx/backends/tabicl.py                | 43 ++++++++++--
 src/tabctx/backends/tabicl_compile.py        | 70 ++++++++++++++++++++
 src/tabctx/serve/factory.py                  | 20 ++++--
 tests/unit/test_serve_factory.py             | 11 +++
 tests/unit/test_tabicl_thread_local_cache.py | 43 ++++++++++++
 9 files changed, 211 insertions(+), 18 deletions(-)
 create mode 100644 src/tabctx/backends/tabicl_compile.py

diff --git a/benchmarks/bench_compile.py b/benchmarks/bench_compile.py
index b078658..85af10f 100644
--- a/benchmarks/bench_compile.py
+++ b/benchmarks/bench_compile.py
@@ -45,6 +45,24 @@ def make_table(rows: int, features: int, seed: int):
     return X, y
 
 
+def symbolic_logn() -> None:
+    """tabicl's ssmax scales queries by log(n) computed with math.log on a
+    Python int, which specializes every compiled graph to one row count
+    (each new table recompiled for 11 s). Inside compiled code, compute it
+    as a tensor op on the symbolic n instead."""
+    from tabicl._model import ssmax
+
+    eager = ssmax._logn
+
+    def logn(n, device, dtype):
+        if torch.compiler.is_compiling():
+            ones = torch.ones((), device=device, dtype=torch.float32)
+            return (ones * torch.sym_max(n, 1)).log().to(dtype)
+        return eager(n, device, dtype)
+
+    ssmax._logn = logn
+
+
 def install(variant: str, mode: str) -> None:
     """Patch tabicl's attention block class (all instances share one graph:
     nn-module parameters are graph inputs, so instances don't recompile)."""
@@ -56,6 +74,7 @@ def install(variant: str, mode: str) -> None:
         # Blocks see several static shapes (column vs row vs ICL stages,
         # induced points vs features): allow a graph per shape family.
         torch._dynamo.config.recompile_limit = 64
+        symbolic_logn()
         Block.forward = torch.compile(Block.forward, dynamic=True, mode=mode)
         return
     raise SystemExit(f"unknown variant {variant}")
diff --git a/benchmarks/bench_overlap.py b/benchmarks/bench_overlap.py
index aae2661..885e59e 100644
--- a/benchmarks/bench_overlap.py
+++ b/benchmarks/bench_overlap.py
@@ -55,10 +55,13 @@ def main() -> int:
     p.add_argument("--context", type=parse_shape, default=(2_000, 20))
     p.add_argument("--test-rows", type=int, default=100)
     p.add_argument("--out", default=None)
+    p.add_argument("--compile", action="store_true", help="compiled TabICL blocks")
     args = p.parse_args()
     device = torch.device(args.device)
 
-    backend = build_backend(args.backend, args.device, "kv")
+    backend = build_backend(
+        args.backend, args.device, "kv", compile_blocks=args.compile
+    )
     backend.warmup()
     cr, cf = args.context
     Xc, yc = make_table(cr, cf, "classification", seed=1)
diff --git a/benchmarks/eval_kv_precision.py b/benchmarks/eval_kv_precision.py
index 986c23e..cb9f2b6 100644
--- a/benchmarks/eval_kv_precision.py
+++ b/benchmarks/eval_kv_precision.py
@@ -11,6 +11,10 @@
 synthetic tables, where KV size actually matters.
 
 Usage: eval_kv_precision.py --device cuda:0 --out kv-precision-a100.json
+
+With --compile, the same with compiled attention blocks (#32/#36): run it
+once with and once without and compare the metrics (compilation patches
+tabicl process-wide, so eager and compiled can't share a process).
 """
 
 from __future__ import annotations
@@ -79,11 +83,15 @@ def main() -> int:
     p = argparse.ArgumentParser()
     p.add_argument("--device", default="cuda:0")
     p.add_argument("--out", default=None)
+    p.add_argument("--precisions", nargs="+", default=list(PRECISIONS))
+    p.add_argument("--compile", action="store_true", help="compiled TabICL blocks")
     args = p.parse_args()
 
     backends = {
-        prec: TabICLBackend(device=args.device, kv_precision=prec)
-        for prec in PRECISIONS
+        prec: TabICLBackend(
+            device=args.device, kv_precision=prec, compile_blocks=args.compile
+        )
+        for prec in args.precisions
     }
     for b in backends.values():
         b.warmup()
diff --git a/benchmarks/profile_fit_predict.py b/benchmarks/profile_fit_predict.py
index 52a54a5..e7f4d2e 100644
--- a/benchmarks/profile_fit_predict.py
+++ b/benchmarks/profile_fit_predict.py
@@ -50,12 +50,14 @@ def make_table(rows: int, features: int, task: str, seed: int):
     return X, y
 
 
-def build_backend(kind: str, device: str, cache_mode: str):
+def build_backend(
+    kind: str, device: str, cache_mode: str, compile_blocks: bool = False
+):
     if kind == "tabicl":
         from tabctx.backends.tabicl import TabICLBackend
 
         kv: bool | str = False if cache_mode == "off" else cache_mode
-        return TabICLBackend(device=device, kv_cache=kv)
+        return TabICLBackend(device=device, kv_cache=kv, compile_blocks=compile_blocks)
     from tabctx.backends.tabpfn import TabPFNBackend
 
     return TabPFNBackend(device=device, cache_mode=cache_mode)
diff --git a/src/tabctx/backends/tabicl.py b/src/tabctx/backends/tabicl.py
index 2f79877..c008ed9 100644
--- a/src/tabctx/backends/tabicl.py
+++ b/src/tabctx/backends/tabicl.py
@@ -147,8 +147,20 @@ def install_cached_logn() -> bool:
     def cached(n: int, device: Any, dtype: Any) -> Any:
         return original(n, device, dtype)
 
-    cached._tabctx_cached = True  # type: ignore[attr-defined]
-    ssmax._logn = cached
+    def logn(n: Any, device: Any, dtype: Any) -> Any:
+        import torch
+
+        if torch.compiler.is_compiling():
+            # Inside a compiled block (tabicl_compile.py) n is symbolic:
+            # math.log would specialize the graph to one row count, so
+            # compute log(n) as a tensor op (in fp32, then cast, as the
+            # eager path rounds a float64 log to dtype).
+            ones = torch.ones((), device=device, dtype=torch.float32)
+            return (ones * torch.sym_max(n, 1)).log().to(dtype)
+        return cached(n, device, dtype)
+
+    logn._tabctx_cached = True  # type: ignore[attr-defined]
+    ssmax._logn = logn
     return True
 
 
@@ -200,6 +212,7 @@ def __init__(
         kv_precision: KVPrecision = "native",
         offload_dir: str | None = None,
         cpu_safety_factor: float = 0.6,
+        compile_blocks: bool = False,
     ) -> None:
         """kv_cache: passed through to TabICL ("kv", "repr", True, or
         False). tabicl's own default is False, which makes every
@@ -211,7 +224,11 @@ def __init__(
         memory-hungry -- the memory cost lands inside fit()'s measured
         delta, so cache accounting and the adaptive admission gate see it
         automatically; "repr" trades ~24x less ICL-cache memory for
-        re-running the ICL stack per predict)."""
+        re-running the ICL stack per predict).
+
+        compile_blocks: torch.compile tabicl's attention blocks (CUDA only;
+        see backends/tabicl_compile.py): 1.7-2.2x faster GPU half, ~20%
+        lower peak, paid for by a one-time compile in warmup()."""
         import torch
 
         self._device = device or ("cuda" if torch.cuda.is_available() else "cpu")
@@ -251,6 +268,11 @@ def __init__(
         _install_restore_hook(TabICLRegressor)
         if not install_cached_logn():
             logger.warning("tabicl.ssmax._logn not found: per-call GPU syncs remain")
+        self._compiled = False
+        if compile_blocks and is_cuda(self._device):
+            from tabctx.backends.tabicl_compile import install_compiled_blocks
+
+            self._compiled = install_compiled_blocks()
 
     def _make_estimator(self, task: Task) -> Any:
         options = {
@@ -491,11 +513,18 @@ def warmup(self) -> None:
         report the weights as part of its context (resident bytes for a
         tiny table were the whole checkpoint), overcharging the cache and
         leaving the books wrong once it is evicted -- and would pay the
-        load latency."""
+        load latency. With compiled blocks, also compile every call family
+        here (fp32 below tabicl's autocast threshold of 1,024 rows, fp16
+        above; fit and cached predict) rather than in tenants' calls."""
         rng = np.random.default_rng(0)
-        X = rng.normal(size=(32, 4))
-        self.fit(X, (X[:, 0] > 0).astype(int), "classification")
-        self.fit(X, X[:, 0], "regression")
+        sizes = (32, 1100) if self._compiled else (32,)
+        for rows in sizes:
+            X = rng.normal(size=(rows, 4))
+            labels = (X[:, 0] > 0).astype(int)
+            for task, y in (("classification", labels), ("regression", X[:, 0])):
+                payload = self.fit(X, y, task)
+                if self._compiled:
+                    self.predict(payload, X[:8], return_proba=task == "classification")
         self._last_context_bytes = None
         self._last_fit_peak_bytes = None
 
diff --git a/src/tabctx/backends/tabicl_compile.py b/src/tabctx/backends/tabicl_compile.py
new file mode 100644
index 0000000..5be0407
--- /dev/null
+++ b/src/tabctx/backends/tabicl_compile.py
@@ -0,0 +1,70 @@
+"""torch.compile for TabICL's attention blocks (#32, #36).
+
+The fit profile (#31) put layer norm at ~21% of the GPU half: memory-bound
+passes over [ensembles, features, rows, 128] activations, each beside a
+residual add or the GELU MLP that streams the same tensor again. Compiling
+each attention block's forward lets Inductor fuse those passes around the
+attention kernel (which stays one kernel). Measured on an A100 at tabicl's
+defaults (fp16 autocast above 1,024 rows), benchmarks/bench_compile.py:
+
+  table       GPU half eager -> compiled   peak GiB
+  3000x20     0.18 -> 0.12 s               2.0 -> 1.7
+  10000x50    0.91 -> 0.52 s               9.9 -> 7.3
+  40000x100   7.65 -> 4.30 s               33.1 -> 27.0
+  20000x200   6.51 -> 2.92 s               31.1 -> 23.4
+
+Accuracy and log loss are unchanged (within 0.0004 log loss on held-out
+rows). In fp32 the compiled blocks match eager to 3e-5 in probability;
+under fp16 autocast they differ by up to ~0.01, the same order as fp16
+itself differs from fp32 (~0.1), from where rounding happens.
+
+Shapes. Every block is compiled with dynamic shapes, so a new table
+doesn't recompile -- once tabicl's ssmax log(n) is computed symbolically
+inside compiled code (see tabicl.install_cached_logn): math.log on a
+Python int specialized every graph to one row count, and each new table
+recompiled for ~11 s. What remains is one graph per call family (stage,
+fp16 vs fp32, cached vs uncached, ...): ~40, compiled by the backend's
+warmup at replica start (~75 s per worker on an A100, workers in
+parallel). Very large tables reach a few more (tabicl batches them
+differently): 1-3 graphs, ~10 s once per worker. Inductor's on-disk cache
+makes later processes faster.
+
+Failure. Compilation errors fall back to eager (dynamo's suppress_errors),
+so a missing compiler toolchain costs speed, not availability.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+# Call families per block (stage x fp16/fp32 x cached/uncached x k given
+# or not): well above dynamo's default of 8, which would silently fall
+# back to eager for the rest.
+_RECOMPILE_LIMIT = 64
+
+
+def install_compiled_blocks(mode: str = "default") -> bool:
+    """Compile tabicl's MultiheadAttentionBlock.forward for every instance
+    (nn-module parameters are graph inputs, so instances share graphs).
+    Idempotent; returns False (and changes nothing) if tabicl or
+    torch.compile isn't available."""
+    try:
+        import torch
+        import torch._dynamo
+        from tabicl._model.layers import MultiheadAttentionBlock
+    except ImportError:
+        return False
+    forward = MultiheadAttentionBlock.forward
+    if getattr(forward, "_tabctx_compiled", False):
+        return True
+    config: Any = torch._dynamo.config
+    config.recompile_limit = max(config.recompile_limit, _RECOMPILE_LIMIT)
+    config.suppress_errors = True
+    compiled = torch.compile(forward, dynamic=True, mode=mode)
+    compiled._tabctx_compiled = True  # type: ignore[attr-defined]
+    MultiheadAttentionBlock.forward = compiled
+    logger.info("tabicl attention blocks compiled (mode=%s)", mode)
+    return True
diff --git a/src/tabctx/serve/factory.py b/src/tabctx/serve/factory.py
index 8166cfa..450a92d 100644
--- a/src/tabctx/serve/factory.py
+++ b/src/tabctx/serve/factory.py
@@ -109,6 +109,7 @@
 UPLOAD_TTL_S_ENV_VAR = "TABCTX_UPLOAD_TTL_S"
 DEVICES_ENV_VAR = "TABCTX_DEVICES"
 WORKER_PROCESSES_ENV_VAR = "TABCTX_WORKER_PROCESSES"
+COMPILE_ENV_VAR = "TABCTX_COMPILE"
 
 logger = logging.getLogger(__name__)
 
@@ -136,6 +137,9 @@ class ServeSettings:
     # Several devices: one worker process per GPU (default) or threads of
     # this process (see _build_process_engine for why processes).
     worker_processes: bool = True
+    # torch.compile the models' hot blocks where measured to help (TabICL's
+    # attention blocks: backends/tabicl_compile.py).
+    compile_models: bool = True
 
     @classmethod
     def from_env(cls) -> ServeSettings:
@@ -212,11 +216,6 @@ def from_env(cls) -> ServeSettings:
             raise ValueError(f"{SPILL_CAPACITY_ENV_VAR} must be positive")
         devices = _parse_devices(os.environ.get(DEVICES_ENV_VAR, "auto"))
         offload_dir = os.environ.get(OFFLOAD_DIR_ENV_VAR) or None
-        raw_workers = os.environ.get(WORKER_PROCESSES_ENV_VAR, "1").strip().lower()
-        if raw_workers not in ("0", "1", "true", "false"):
-            raise ValueError(
-                f"{WORKER_PROCESSES_ENV_VAR}={raw_workers!r} must be 0 or 1"
-            )
         return cls(
             backends=backends,
             gpu_memory_fraction=fraction,
@@ -229,10 +228,18 @@ def from_env(cls) -> ServeSettings:
             spill_capacity_bytes=spill_capacity,
             devices=devices,
             offload_dir=offload_dir,
-            worker_processes=raw_workers in ("1", "true"),
+            worker_processes=_env_flag(WORKER_PROCESSES_ENV_VAR, cls.worker_processes),
+            compile_models=_env_flag(COMPILE_ENV_VAR, cls.compile_models),
         )
 
 
+def _env_flag(name: str, default: bool) -> bool:
+    raw = os.environ.get(name, "1" if default else "0").strip().lower()
+    if raw not in ("0", "1", "true", "false"):
+        raise ValueError(f"{name}={raw!r} must be 0 or 1")
+    return raw in ("1", "true")
+
+
 def _parse_devices(raw: str) -> tuple[str, ...]:
     raw = raw.strip().lower()
     if raw in ("", "auto"):
@@ -499,6 +506,7 @@ def _build_real_backend(
             # Concurrent offloading calls each get their share of host
             # headroom (see memory/host.py).
             cpu_safety_factor=0.6 / host_slots,
+            compile_blocks=settings.compile_models,
         )
     # Weights load at replica start, not inside (and charged to) the first
     # tenant's fit.
diff --git a/tests/unit/test_serve_factory.py b/tests/unit/test_serve_factory.py
index 0d99987..ab7274c 100644
--- a/tests/unit/test_serve_factory.py
+++ b/tests/unit/test_serve_factory.py
@@ -471,3 +471,14 @@ def test_default_and_parsing(self, monkeypatch):
         monkeypatch.setenv("TABCTX_KV_PRECISION", "fp4")
         with pytest.raises(ValueError, match="TABCTX_KV_PRECISION"):
             ServeSettings.from_env()
+
+
+def test_compile_flag_from_env(monkeypatch):
+    from tabctx.serve.factory import COMPILE_ENV_VAR
+
+    assert ServeSettings.from_env().compile_models  # on by default
+    monkeypatch.setenv(COMPILE_ENV_VAR, "0")
+    assert not ServeSettings.from_env().compile_models
+    monkeypatch.setenv(COMPILE_ENV_VAR, "maybe")
+    with pytest.raises(ValueError, match="TABCTX_COMPILE"):
+        ServeSettings.from_env()
diff --git a/tests/unit/test_tabicl_thread_local_cache.py b/tests/unit/test_tabicl_thread_local_cache.py
index 9e88cf1..074348e 100644
--- a/tests/unit/test_tabicl_thread_local_cache.py
+++ b/tests/unit/test_tabicl_thread_local_cache.py
@@ -74,3 +74,46 @@ def test_cached_logn_matches_and_is_installed_once():
             assert torch.equal(patched(n, torch.device("cpu"), torch.float16), expected)
     finally:
         ssmax._logn = original
+
+
+def test_logn_stays_symbolic_inside_compiled_code():
+    """math.log on the row count specialized every compiled graph to one
+    table size (each new table recompiled for ~11 s on an A100)."""
+    ssmax = pytest.importorskip("tabicl._model.ssmax")
+    from torch._dynamo.testing import CompileCounter
+
+    from tabctx.backends.tabicl import install_cached_logn
+
+    original = ssmax._logn
+    try:
+        install_cached_logn()
+        counter = CompileCounter()
+
+        @torch.compile(backend=counter, dynamic=True)
+        def scale(q):
+            return q * ssmax._logn(q.shape[0], q.device, q.dtype)
+
+        for n in (3000, 5000, 7000, 12000):
+            q = torch.ones(n, 4)
+            expected = original(n, torch.device("cpu"), torch.float32)
+            assert torch.allclose(scale(q)[0, 0], expected, rtol=1e-6)
+        assert counter.frame_count == 1  # one graph for every row count
+    finally:
+        ssmax._logn = original
+        torch._dynamo.reset()
+
+
+def test_compiled_blocks_install_once():
+    layers = pytest.importorskip("tabicl._model.layers")
+    from tabctx.backends.tabicl_compile import install_compiled_blocks
+
+    original = layers.MultiheadAttentionBlock.forward
+    try:
+        assert install_compiled_blocks()
+        compiled = layers.MultiheadAttentionBlock.forward
+        assert compiled is not original
+        assert install_compiled_blocks()
+        assert layers.MultiheadAttentionBlock.forward is compiled  # idempotent
+    finally:
+        layers.MultiheadAttentionBlock.forward = original
+        torch._dynamo.reset()

From 2de5ff2ce97398335ca6904e5eaacf7cb85dac04 Mon Sep 17 00:00:00 2001
From: Amrit Krishnan <amrit110@gmail.com>
Date: Sun, 18 Oct 2026 00:11:06 -0400
Subject: [PATCH 29/39] docs: torch.compile results in roadmap, changelog and
 README; bench_compile covers TabPFN

bench_compile now builds the backends with their own compile options
(TabICL compile_blocks, TabPFN-3.5's built-in enable_torch_compile)
instead of patching them itself.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 CHANGELOG.md                |   1 +
 README.md                   |   6 ++
 ROADMAP.md                  |  23 ++++--
 benchmarks/bench_compile.py | 146 +++++++++++++++++++++---------------
 4 files changed, 107 insertions(+), 69 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index fc5552d..fc13652 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -14,6 +14,7 @@
 - ROADMAP: multi-GPU and CUDA-level track (#30-#37), PyTorch-level first.
 
 ### Changed
+- **TabICL's attention blocks are compiled** (#36, `TABCTX_COMPILE`, default on; `backends/tabicl_compile.py`). Inductor fuses the memory-bound layer norms, residual adds and GELU MLP around each attention kernel: on an A100 the fit's GPU half is 1.7-2.2x faster (10k x 50: 0.91 -> 0.52 s; 40k x 100: 7.65 -> 4.30 s; 20k x 200: 6.51 -> 2.92 s) and its peak ~20% lower (40k x 100: 33.1 -> 27.0 GiB); 100-row predicts 20-35% faster. Accuracy, AUC and R^2 are unchanged on the evaluation datasets, and predicts beside a fit stay correct. Ray Serve, 6 mixed users on 4x A100: 5.78 -> 6.99 predicts/s, small-fit p95 1.48 -> 0.95 s. tabicl's ssmax log(n) is now computed symbolically inside compiled code, so a new table size doesn't recompile. The graphs (~40) compile in `warmup()` at replica start (~75 s per worker, workers in parallel); compile errors fall back to eager.
 - **Host memory is admitted by estimated bytes, node-wide** (`memory/host.py`). Each backend estimates a call's host use (TabICL ~5 KB per cell when it offloads stage outputs, else nothing; TabPFN ~0.6 KB per cell) and the guard admits calls against one budget shared by every worker process, first come first served (a large fit can't be starved by modest ones), with at most `slots` offloading calls at once (each budgets 1/slots of free host memory for staging). Calls under 1 GiB never wait. The old guard gave every call over 3M cells a whole 48 GiB slot, which serialized fits that use a few GiB (40k x 100 TabICL fits took p50 195 s under concurrent load on the 4x L4 rehearsal). A worker that dies holding host memory has it returned by its parent, and a stuck accounting lock fails calls with a retryable 503 instead of hanging. **`/readyz` `host_memory` keys changed**: `large_fits_active`/`large_fits_waiting` are now `calls_active`/`calls_waiting`, plus `reserved_bytes` and `calls_offloading`.
 - **TabICL fits prepare their CPU half before taking the device**: preprocessing (~30-37% of a fit) overlaps another fit's GPU half -- back-to-back fits 17-26% faster on an A100. At most one fit per device prepares ahead, so queued fits don't each hold preprocessed copies in host memory.
 - Worker processes shut down together under one 10 s deadline (was up to 10 s each, past a pod's grace period).
diff --git a/README.md b/README.md
index 8e3bb77..27e09ea 100644
--- a/README.md
+++ b/README.md
@@ -389,6 +389,12 @@ pip install "tabctx[dev]"        # + test dependencies
   (`TABCTX_KV_CACHE=kv|repr|off`), with predictions verified identical
   to the uncached path -- and loads the pretrained backbone once per
   process instead of per fit.
+- **Compiled attention blocks** (`TABCTX_COMPILE=1`, the default):
+  TabICL's attention blocks run through `torch.compile`, which fuses the
+  memory-bound layer norms and residual adds around each attention
+  kernel. Fits are 1.7-2.2x faster on the GPU with ~20% lower peak
+  memory, and accuracy is unchanged. The compile happens once per worker
+  at startup (~75 s), not in tenants' requests.
 - **Same-context coalescing** (`CoalescingPredictor`): concurrent
   requests against one cached context are packed into a single GPU call
   within a ~5ms window and split back per caller, amortizing per-call
diff --git a/ROADMAP.md b/ROADMAP.md
index de24322..cf6c0fc 100644
--- a/ROADMAP.md
+++ b/ROADMAP.md
@@ -71,10 +71,11 @@ no FP8).
    half is prepared before taking the device (overlaps another fit's GPU
    half: back-to-back fits 17-26% faster), and tabicl's per-call
    log(n) host-to-device syncs are cached (bit-identical; 10k x 50 fits
-   16% faster). Next from it: fuse the column stage's layer norms
-   (`torch.compile`, #32/#36), and a GPU Yeo-Johnson fit for isolated-fit
-   latency (~3 ms vs 2.9 s at 40k x 100; not bit-identical to sklearn's
-   lambdas, so behind an accuracy study).
+   16% faster), and the attention blocks are compiled (#36, below). Next
+   from it: a GPU Yeo-Johnson fit for isolated-fit latency (~3 ms vs
+   2.9 s at 40k x 100; not bit-identical to sklearn's lambdas, so behind
+   an accuracy study). TabPFN-3.5's fit is the same shape: RMS norm 33% of
+   its CUDA time, flash attention 26%, copies 16%.
 7. **Per-device calibration grids** ([#37](https://github.com/VectorInstitute/tabctx/issues/37), TabPFN done): TabPFN-3.5 has its
    own requirement grid (A100, 20 shapes): it, too, sizes working memory
    from what is free, so admission reaches 100k rows at 100 features on an
@@ -95,9 +96,17 @@ no FP8).
 12. **Context parallelism across GPUs** ([#35](https://github.com/VectorInstitute/tabctx/issues/35)): tables larger than one GPU
     (ring attention over the replica's devices). Disk offload already
     lets one A100 fit 130k x 500 (65M cells) in 26 min.
-13. **CUDA graphs / torch.compile** ([#36](https://github.com/VectorInstitute/tabctx/issues/36)): the GPU half is now mostly GPU-
-    bound, so the remaining gains are fusion (layer norm, elementwise)
-    rather than launch overhead.
+13. **torch.compile** ([#36](https://github.com/VectorInstitute/tabctx/issues/36), done for TabICL): each attention block's
+    forward compiled, so Inductor fuses the layer norms, residual adds and
+    GELU MLP around the attention kernel. A100: GPU half 1.7-2.2x faster
+    (40k x 100: 7.65 -> 4.30 s), peak memory ~20% lower (33.1 -> 27.0
+    GiB, so larger tables fit), accuracy unchanged on the evaluation sets;
+    Ray Serve, 6 users on 4x A100: 5.78 -> 6.99 predicts/s. Dynamic
+    shapes needed tabicl's ssmax log(n) computed symbolically (math.log on
+    the row count recompiled every new table). ~40 graphs compile in
+    warmup (~75 s per worker). CUDA graphs aren't worth it: the GPU half
+    is GPU-bound, not launch-bound. TabPFN-3.5's own compile path is
+    being measured.
 14. **Serving-framework-agnostic core** ([#39](https://github.com/VectorInstitute/tabctx/issues/39)): the router, per-GPU workers
     and scheduling are plain Python with no Ray dependency; a Dynamo
     worker hosts the same core. Next: residency events and pluggable
diff --git a/benchmarks/bench_compile.py b/benchmarks/bench_compile.py
index 85af10f..dfc83f2 100644
--- a/benchmarks/bench_compile.py
+++ b/benchmarks/bench_compile.py
@@ -1,22 +1,24 @@
 #!/usr/bin/env python3
-"""Does torch.compile speed up TabICL's GPU half, and is it exact enough? (#32/#36)
-
-The fit profile (#31) put layer norm at ~21% of the GPU half -- memory-bound
-passes over [ensembles, features, rows, 128] activations, each next to a
-residual add or a GELU MLP that also streams the tensor. Inductor can fuse
-those elementwise passes. This compares, on one device:
-
-- eager (baseline);
-- ``block``: each attention block's whole forward compiled (norms,
-  residual adds and the GELU MLP fused around the attention kernel).
-
-For each it reports the GPU half's time (TabICLBackend.fit_prepared, CPU
-preprocessing excluded), the first call's compile cost, predict latency, and
-the largest probability difference against eager on held-out rows.
-
-Usage (on a GPU host with tabicl and the weights):
-  bench_compile.py --device cuda:0 --shapes 10000x50 40000x100 --repeats 3 \
-      --variant block --baseline eager.npz --out compile-a100.jsonl
+"""Does torch.compile speed up a backend's fit and predict, and is it exact
+enough? (#32/#36)
+
+Compares one backend eager vs compiled, on one device (compilation patches
+the model process-wide, so run eager and compiled as separate processes):
+
+- TabICL: its attention blocks compiled (TabICLBackend(compile_blocks=True),
+  see backends/tabicl_compile.py) -- layer norms, residual adds and the
+  GELU MLP fused around the attention kernel;
+- TabPFN-3.5: the library's own compile path (enable_torch_compile).
+
+Reports the fit time (TabICL: its GPU half, CPU preprocessing excluded),
+the first call's compile cost, predict latency, accuracy and log loss on
+held-out rows, and the largest probability difference against a saved
+eager run.
+
+Usage (on a GPU host with the models' weights):
+  bench_compile.py --backend tabicl --shapes 10000x50 40000x100 --save-proba eager.npz
+  bench_compile.py --backend tabicl --shapes 10000x50 40000x100 --compile \
+      --baseline eager.npz --out compile-a100.jsonl
 """
 
 from __future__ import annotations
@@ -45,46 +47,64 @@ def make_table(rows: int, features: int, seed: int):
     return X, y
 
 
-def symbolic_logn() -> None:
-    """tabicl's ssmax scales queries by log(n) computed with math.log on a
-    Python int, which specializes every compiled graph to one row count
-    (each new table recompiled for 11 s). Inside compiled code, compute it
-    as a tensor op on the symbolic n instead."""
-    from tabicl._model import ssmax
-
-    eager = ssmax._logn
+def enable_tabpfn_compile() -> None:
+    """TabPFN-3.5 ships its own torch.compile path (PerformanceOptions.
+    enable_torch_compile: selected sub-graphs, dynamic row/feature dims,
+    symbolic log n), off by default and not exposed by the estimators;
+    turn it on in the model's default options."""
+    import dataclasses
 
-    def logn(n, device, dtype):
-        if torch.compiler.is_compiling():
-            ones = torch.ones((), device=device, dtype=torch.float32)
-            return (ones * torch.sym_max(n, 1)).log().to(dtype)
-        return eager(n, device, dtype)
+    from tabpfn.architectures.tabpfn_v3_5 import TabPFNV3p5
 
-    ssmax._logn = logn
+    defaults = TabPFNV3p5.get_default_performance_options
 
+    def with_compile(self):
+        return dataclasses.replace(defaults(self), enable_torch_compile=True)
 
-def install(variant: str, mode: str) -> None:
-    """Patch tabicl's attention block class (all instances share one graph:
-    nn-module parameters are graph inputs, so instances don't recompile)."""
-    from tabicl._model.layers import MultiheadAttentionBlock as Block
+    TabPFNV3p5.get_default_performance_options = with_compile
 
-    if variant == "eager":
-        return
-    if variant == "block":
-        # Blocks see several static shapes (column vs row vs ICL stages,
-        # induced points vs features): allow a graph per shape family.
-        torch._dynamo.config.recompile_limit = 64
-        symbolic_logn()
-        Block.forward = torch.compile(Block.forward, dynamic=True, mode=mode)
-        return
-    raise SystemExit(f"unknown variant {variant}")
 
+def build(kind: str, device: str, compiled: bool):
+    if kind == "tabpfn":
+        from tabctx.backends.tabpfn import TabPFNBackend
 
-def run(device: str, rows: int, features: int, repeats: int, amp: str) -> dict:
+        if compiled:
+            enable_tabpfn_compile()
+        return TabPFNBackend(device=device)
     from tabctx.backends.tabicl import TabICLBackend
 
-    backend = TabICLBackend(device=device)
-    if amp != "auto":  # tabicl's default: fp16 autocast above 1024 rows
+    return TabICLBackend(device=device, compile_blocks=compiled)
+
+
+def fit_timed(backend, X, y, device: str):
+    """The fit's GPU half where the backend splits it (TabICL), else all."""
+    prepare = getattr(backend, "prepare_fit", None)
+    if callable(prepare):
+        prepared = prepare(X, y, "classification")
+        torch.cuda.synchronize(device)
+        t0 = time.perf_counter()
+        payload = backend.fit_prepared(prepared)
+    else:
+        torch.cuda.synchronize(device)
+        t0 = time.perf_counter()
+        payload = backend.fit(X, y, "classification")
+    torch.cuda.synchronize(device)
+    return payload, time.perf_counter() - t0
+
+
+def run(
+    kind: str,
+    device: str,
+    compiled: bool,
+    rows: int,
+    features: int,
+    repeats: int,
+    amp: str,
+) -> dict:
+    backend = build(kind, device, compiled)
+    if (
+        amp != "auto" and kind == "tabicl"
+    ):  # tabicl's default: fp16 autocast above 1024 rows
         make = backend._make_estimator
 
         def make_with_amp(task):
@@ -100,12 +120,7 @@ def make_with_amp(task):
     times, payload = [], None
     first = 0.0
     for i in range(repeats + 1):
-        prepared = backend.prepare_fit(X, y, "classification")
-        torch.cuda.synchronize(device)
-        t0 = time.perf_counter()
-        payload = backend.fit_prepared(prepared)
-        torch.cuda.synchronize(device)
-        elapsed = time.perf_counter() - t0
+        payload, elapsed = fit_timed(backend, X, y, device)
         if i == 0:
             first = elapsed
         else:
@@ -121,7 +136,7 @@ def make_with_amp(task):
         "accuracy": round(float((proba.argmax(1) == y_test).mean()), 4),
         "log_loss": round(float(-np.log(p_true).mean()), 5),
         "first_s": round(first, 3),
-        "gpu_half_s": round(float(np.median(times)), 3),
+        "fit_s": round(float(np.median(times)), 3),
         "predict_s": round(time.perf_counter() - t0, 3),
         "proba": proba,
         "peak_gib": round(torch.cuda.max_memory_allocated(device) / 1024**3, 2),
@@ -133,20 +148,27 @@ def main() -> int:
     ap.add_argument("--device", default="cuda:0")
     ap.add_argument("--shapes", nargs="+", default=["10000x50", "40000x100"])
     ap.add_argument("--repeats", type=int, default=3)
-    ap.add_argument("--variant", default="eager")
-    ap.add_argument("--mode", default="default")
+    ap.add_argument("--backend", choices=("tabicl", "tabpfn"), default="tabicl")
+    ap.add_argument("--compile", action="store_true")
     ap.add_argument("--amp", choices=("auto", "on", "off"), default="auto")
     ap.add_argument("--baseline", help="npz of eager probabilities to compare against")
     ap.add_argument("--save-proba", help="write this run's probabilities (npz)")
     ap.add_argument("--out")
     args = ap.parse_args()
 
-    install(args.variant, args.mode)
     baseline = dict(np.load(args.baseline)) if args.baseline else {}
     saved = {}
     for text in args.shapes:
         rows, features = parse_shape(text)
-        result = run(args.device, rows, features, args.repeats, args.amp)
+        result = run(
+            args.backend,
+            args.device,
+            args.compile,
+            rows,
+            features,
+            args.repeats,
+            args.amp,
+        )
         proba = result.pop("proba")
         saved[text] = proba
         if text in baseline:
@@ -154,8 +176,8 @@ def main() -> int:
             result["max_abs_proba_diff"] = float(np.abs(proba - ref).max())
             result["label_agreement"] = float((proba.argmax(1) == ref.argmax(1)).mean())
         record = {
-            "variant": args.variant,
-            "mode": args.mode,
+            "backend": args.backend,
+            "compiled": args.compile,
             "amp": args.amp,
             "shape": text,
             **result,

From 950e5ca1f109fa2aa4cc960a871cc561f758b03b Mon Sep 17 00:00:00 2001
From: Amrit Krishnan <amrit110@gmail.com>
Date: Sun, 18 Oct 2026 00:42:57 -0400
Subject: [PATCH 30/39] fix: review of the host guard and compile changes

Host guard (memory/host.py):
- Across processes it no longer uses mp.Condition: a process killed in
  wait() left its sleeper count behind, so the next notify_all blocked
  forever holding the lock (reproduced: release_owner hung after a waiter
  was killed). A plain lock that waiters poll (20 ms) instead; in-process
  keeps a condition variable.
- A call's release never raises: a stuck lock is logged instead of turning
  a successful fit into a 503 (and leaking its reservation). A waiter that
  loses the lock still unparks its ticket.
- Queue slots are sized per device for the in-process guard.

Engine: the prepare gate only limits large fits (1M+ cells) and is
returned once the fit has the device, so large fits waiting for host
memory no longer hold up small fits. host_bytes_estimate without the
`offloads` keyword still works.

Compile: only tabicl's fp16 autocast path (1,024+ rows) is compiled
(small tables gain little and doubled the graphs); warmup covers 1-row
predicts and 1-feature tables (~50 graphs, ~130 s per worker; 0 compiles
in requests for common shapes on an A100); the recompile limit was being
hit (64 at exactly 64 graphs) and is now 256 per device; a warning if
nothing compiled; warmup resets every measurement.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 CHANGELOG.md                                 |   2 +-
 README.md                                    |   5 +-
 ROADMAP.md                                   |  11 +-
 benchmarks/bench_compile.py                  |   6 +
 src/tabctx/backends/tabicl.py                |  26 ++--
 src/tabctx/backends/tabicl_compile.py        |  60 ++++++++--
 src/tabctx/engine.py                         |  26 ++--
 src/tabctx/memory/host.py                    | 120 +++++++++++++------
 src/tabctx/serve/factory.py                  |  15 ++-
 tests/unit/test_engine_scheduling.py         |  28 ++++-
 tests/unit/test_host_memory.py               |  70 +++++++++++
 tests/unit/test_tabicl_thread_local_cache.py |   1 +
 12 files changed, 296 insertions(+), 74 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index fc13652..7ace6c5 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -14,7 +14,7 @@
 - ROADMAP: multi-GPU and CUDA-level track (#30-#37), PyTorch-level first.
 
 ### Changed
-- **TabICL's attention blocks are compiled** (#36, `TABCTX_COMPILE`, default on; `backends/tabicl_compile.py`). Inductor fuses the memory-bound layer norms, residual adds and GELU MLP around each attention kernel: on an A100 the fit's GPU half is 1.7-2.2x faster (10k x 50: 0.91 -> 0.52 s; 40k x 100: 7.65 -> 4.30 s; 20k x 200: 6.51 -> 2.92 s) and its peak ~20% lower (40k x 100: 33.1 -> 27.0 GiB); 100-row predicts 20-35% faster. Accuracy, AUC and R^2 are unchanged on the evaluation datasets, and predicts beside a fit stay correct. Ray Serve, 6 mixed users on 4x A100: 5.78 -> 6.99 predicts/s, small-fit p95 1.48 -> 0.95 s. tabicl's ssmax log(n) is now computed symbolically inside compiled code, so a new table size doesn't recompile. The graphs (~40) compile in `warmup()` at replica start (~75 s per worker, workers in parallel); compile errors fall back to eager.
+- **TabICL's attention blocks are compiled** (#36, `TABCTX_COMPILE`, default on; `backends/tabicl_compile.py`). Inductor fuses the memory-bound layer norms, residual adds and GELU MLP around each attention kernel: on an A100 the fit's GPU half is 1.7-2.2x faster (10k x 50: 0.91 -> 0.52 s; 40k x 100: 7.65 -> 4.30 s; 20k x 200: 6.51 -> 2.92 s) and its peak ~20% lower (40k x 100: 33.1 -> 27.0 GiB); 100-row predicts 20-35% faster. Accuracy, AUC and R^2 are unchanged on the evaluation datasets, and predicts beside a fit stay correct. Ray Serve, 6 mixed users on 4x A100: 5.78 -> 6.99 predicts/s, small-fit p95 1.48 -> 0.95 s. tabicl's ssmax log(n) is now computed symbolically inside compiled code, so a new table size doesn't recompile. Only tabicl's fp16 path (1,024+ rows) is compiled; its common graphs (~50) compile in `warmup()` at replica start (~130 s per worker, workers in parallel); shapes it doesn't cover (very large tables) compile on first use, ~10 s once per worker. Compile errors fall back to eager, with a warning after warmup if nothing compiled.
 - **Host memory is admitted by estimated bytes, node-wide** (`memory/host.py`). Each backend estimates a call's host use (TabICL ~5 KB per cell when it offloads stage outputs, else nothing; TabPFN ~0.6 KB per cell) and the guard admits calls against one budget shared by every worker process, first come first served (a large fit can't be starved by modest ones), with at most `slots` offloading calls at once (each budgets 1/slots of free host memory for staging). Calls under 1 GiB never wait. The old guard gave every call over 3M cells a whole 48 GiB slot, which serialized fits that use a few GiB (40k x 100 TabICL fits took p50 195 s under concurrent load on the 4x L4 rehearsal). A worker that dies holding host memory has it returned by its parent, and a stuck accounting lock fails calls with a retryable 503 instead of hanging. **`/readyz` `host_memory` keys changed**: `large_fits_active`/`large_fits_waiting` are now `calls_active`/`calls_waiting`, plus `reserved_bytes` and `calls_offloading`.
 - **TabICL fits prepare their CPU half before taking the device**: preprocessing (~30-37% of a fit) overlaps another fit's GPU half -- back-to-back fits 17-26% faster on an A100. At most one fit per device prepares ahead, so queued fits don't each hold preprocessed copies in host memory.
 - Worker processes shut down together under one 10 s deadline (was up to 10 s each, past a pod's grace period).
diff --git a/README.md b/README.md
index 27e09ea..f989823 100644
--- a/README.md
+++ b/README.md
@@ -393,8 +393,9 @@ pip install "tabctx[dev]"        # + test dependencies
   TabICL's attention blocks run through `torch.compile`, which fuses the
   memory-bound layer norms and residual adds around each attention
   kernel. Fits are 1.7-2.2x faster on the GPU with ~20% lower peak
-  memory, and accuracy is unchanged. The compile happens once per worker
-  at startup (~75 s), not in tenants' requests.
+  memory, and accuracy is unchanged. The compile happens mostly once per
+  worker at startup (~130 s); a call shape the startup didn't cover (very
+  large tables, for one) compiles on first use, ~10 s once per worker.
 - **Same-context coalescing** (`CoalescingPredictor`): concurrent
   requests against one cached context are packed into a single GPU call
   within a ~5ms window and split back per caller, amortizing per-call
diff --git a/ROADMAP.md b/ROADMAP.md
index cf6c0fc..578fe47 100644
--- a/ROADMAP.md
+++ b/ROADMAP.md
@@ -103,10 +103,13 @@ no FP8).
     GiB, so larger tables fit), accuracy unchanged on the evaluation sets;
     Ray Serve, 6 users on 4x A100: 5.78 -> 6.99 predicts/s. Dynamic
     shapes needed tabicl's ssmax log(n) computed symbolically (math.log on
-    the row count recompiled every new table). ~40 graphs compile in
-    warmup (~75 s per worker). CUDA graphs aren't worth it: the GPU half
-    is GPU-bound, not launch-bound. TabPFN-3.5's own compile path is
-    being measured.
+    the row count recompiled every new table). ~50 graphs (fp16 path
+    only) compile in warmup (~130 s per worker). CUDA graphs aren't worth it: the GPU half
+    is GPU-bound, not launch-bound. TabPFN-3.5's own compile path
+    (`enable_torch_compile`) measured 1.3-1.5x faster fits, ~30% faster
+    predicts and ~22% lower peak, but it still recompiles for many new
+    table sizes (9-82 s each, even with its chunked-SDPA loop kept out of
+    the graph), so it stays off until those are traced down.
 14. **Serving-framework-agnostic core** ([#39](https://github.com/VectorInstitute/tabctx/issues/39)): the router, per-GPU workers
     and scheduling are plain Python with no Ray dependency; a Dynamo
     worker hosts the same core. Next: residency events and pluggable
diff --git a/benchmarks/bench_compile.py b/benchmarks/bench_compile.py
index dfc83f2..19d0cf0 100644
--- a/benchmarks/bench_compile.py
+++ b/benchmarks/bench_compile.py
@@ -62,6 +62,12 @@ def with_compile(self):
         return dataclasses.replace(defaults(self), enable_torch_compile=True)
 
     TabPFNV3p5.get_default_performance_options = with_compile
+    # Its SDPA wrapper loops over ceil(batch*heads / 65536) chunks: a Python
+    # loop count that depends on the table, so every new size recompiled
+    # (60-266 s each). Run it eagerly (a graph break around attention).
+    from tabpfn.architectures.shared import scaled_dot_product_attention as sdpa
+
+    sdpa._torch_sdpa = torch.compiler.disable(sdpa._torch_sdpa)
 
 
 def build(kind: str, device: str, compiled: bool):
diff --git a/src/tabctx/backends/tabicl.py b/src/tabctx/backends/tabicl.py
index c008ed9..e2ec377 100644
--- a/src/tabctx/backends/tabicl.py
+++ b/src/tabctx/backends/tabicl.py
@@ -513,20 +513,32 @@ def warmup(self) -> None:
         report the weights as part of its context (resident bytes for a
         tiny table were the whole checkpoint), overcharging the cache and
         leaving the books wrong once it is evicted -- and would pay the
-        load latency. With compiled blocks, also compile every call family
-        here (fp32 below tabicl's autocast threshold of 1,024 rows, fp16
-        above; fit and cached predict) rather than in tenants' calls."""
+        load latency. With compiled blocks, also compile the common call
+        families here (fp32 and fp16 paths, fits and cached predicts,
+        1-feature tables and 1-row predicts) rather than in tenants'
+        calls."""
         rng = np.random.default_rng(0)
-        sizes = (32, 1100) if self._compiled else (32,)
-        for rows in sizes:
-            X = rng.normal(size=(rows, 4))
+        # (rows, features): with compiled blocks, tabicl's fp16 autocast path
+        # (1,024+ rows; the only one compiled), including a 1-feature table
+        # (dynamo specializes size-1 dimensions, so they get their own
+        # graphs).
+        shapes = ((32, 4), (1100, 4), (1100, 1)) if self._compiled else ((32, 4),)
+        for rows, features in shapes:
+            X = rng.normal(size=(rows, features))
             labels = (X[:, 0] > 0).astype(int)
             for task, y in (("classification", labels), ("regression", X[:, 0])):
                 payload = self.fit(X, y, task)
                 if self._compiled:
-                    self.predict(payload, X[:8], return_proba=task == "classification")
+                    proba = task == "classification"
+                    for n in (8, 1):  # 1-row predicts: also size-1 graphs
+                        self.predict(payload, X[:n], return_proba=proba)
+        if self._compiled:
+            from tabctx.backends.tabicl_compile import warn_if_not_compiled
+
+            warn_if_not_compiled()
         self._last_context_bytes = None
         self._last_fit_peak_bytes = None
+        self._last_fit_resident_bytes = None
 
     def context_bytes_hint(self, n_train: int, n_features: int) -> int | None:
         del n_train, n_features
diff --git a/src/tabctx/backends/tabicl_compile.py b/src/tabctx/backends/tabicl_compile.py
index 5be0407..5ea43b9 100644
--- a/src/tabctx/backends/tabicl_compile.py
+++ b/src/tabctx/backends/tabicl_compile.py
@@ -23,11 +23,13 @@
 inside compiled code (see tabicl.install_cached_logn): math.log on a
 Python int specialized every graph to one row count, and each new table
 recompiled for ~11 s. What remains is one graph per call family (stage,
-fp16 vs fp32, cached vs uncached, ...): ~40, compiled by the backend's
-warmup at replica start (~75 s per worker on an A100, workers in
-parallel). Very large tables reach a few more (tabicl batches them
-differently): 1-3 graphs, ~10 s once per worker. Inductor's on-disk cache
-makes later processes faster.
+cached vs uncached, size-1 dimensions, ...). Only tabicl's fp16 autocast
+path (tables of 1,024+ rows) is compiled: smaller tables run in fp32 in
+~0.1 s and would double the graphs. ~50 graphs compile in the backend's
+warmup at replica start (~130 s per worker on an A100, workers in
+parallel; Inductor's cache barely shortens it, tracing dominates). Calls warmup doesn't cover still compile inside a request,
+once per worker: very large tables reach 1-3 more graphs (tabicl batches
+them differently), ~10 s each, and other unusual shapes may too.
 
 Failure. Compilation errors fall back to eager (dynamo's suppress_errors),
 so a missing compiler toolchain costs speed, not availability.
@@ -40,10 +42,11 @@
 
 logger = logging.getLogger(__name__)
 
-# Call families per block (stage x fp16/fp32 x cached/uncached x k given
-# or not): well above dynamo's default of 8, which would silently fall
-# back to eager for the rest.
-_RECOMPILE_LIMIT = 64
+# Graphs per block, per device (stage x fp16/fp32 x cached/uncached x ...,
+# ~50 measured; guards include the device, so an in-process engine over
+# several GPUs needs that many per GPU): well above dynamo's default of 8,
+# past which it silently runs the rest eagerly.
+_RECOMPILE_LIMIT_PER_DEVICE = 256
 
 
 def install_compiled_blocks(mode: str = "default") -> bool:
@@ -60,11 +63,44 @@ def install_compiled_blocks(mode: str = "default") -> bool:
     forward = MultiheadAttentionBlock.forward
     if getattr(forward, "_tabctx_compiled", False):
         return True
+    # Process-wide settings: they apply to any other torch.compile user in
+    # this process too (suppress_errors makes every compile failure fall
+    # back to eager; warn_if_not_compiled reports it for these blocks).
     config: Any = torch._dynamo.config
-    config.recompile_limit = max(config.recompile_limit, _RECOMPILE_LIMIT)
+    devices = max(1, torch.cuda.device_count())
+    config.recompile_limit = max(
+        config.recompile_limit, _RECOMPILE_LIMIT_PER_DEVICE * devices
+    )
     config.suppress_errors = True
     compiled = torch.compile(forward, dynamic=True, mode=mode)
-    compiled._tabctx_compiled = True  # type: ignore[attr-defined]
-    MultiheadAttentionBlock.forward = compiled
+
+    # Only under autocast: tabicl turns fp16 autocast on for tables of
+    # 1,024+ rows, where the fused kernels pay off; smaller tables (fp32,
+    # ~0.1 s fits) gain little and would double the graphs to compile.
+    def dispatch(self: Any, *args: Any, **kwargs: Any) -> Any:
+        if torch.is_autocast_enabled("cuda"):
+            return compiled(self, *args, **kwargs)
+        return forward(self, *args, **kwargs)
+
+    dispatch._tabctx_compiled = True  # type: ignore[attr-defined]
+    dispatch._tabctx_eager = forward  # type: ignore[attr-defined]
+    MultiheadAttentionBlock.forward = dispatch
     logger.info("tabicl attention blocks compiled (mode=%s)", mode)
     return True
+
+
+def warn_if_not_compiled() -> bool:
+    """After warmup: warn if no graph was compiled (compile errors fall back
+    to eager silently, losing the speedup) or a block hit dynamo's limit.
+    Returns whether compilation looks healthy."""
+    from torch._dynamo.utils import counters
+
+    graphs = counters["stats"]["unique_graphs"]
+    if graphs == 0:
+        logger.warning(
+            "tabicl attention blocks did not compile (running eagerly): check "
+            "the worker log for dynamo errors, or set TABCTX_COMPILE=0"
+        )
+        return False
+    logger.info("tabicl attention blocks: %d graphs compiled", graphs)
+    return True
diff --git a/src/tabctx/engine.py b/src/tabctx/engine.py
index 46a5949..6c9120f 100644
--- a/src/tabctx/engine.py
+++ b/src/tabctx/engine.py
@@ -36,8 +36,10 @@
 # A predict whose context is evicted again (under memory pressure) between
 # its restore and getting the device tries this many times.
 _PREDICT_ATTEMPTS = 3
-# Fits prepared (CPU half done) ahead of the one holding the device.
+# Large fits prepared (CPU half done) and waiting for the device at once;
+# smaller fits' prepared copies are too small to limit.
 _PREPARED_AHEAD = 1
+GATED_FIT_CELLS = 1_000_000
 
 
 def _not_found(dataset_id: str) -> DatasetNotFoundError:
@@ -186,7 +188,7 @@ def __init__(
         # _fit_lock, _predict_lock (the cache lock and _release_lock are
         # leaves).
         self._fit_lock = threading.Lock()
-        self._prepare_gate = threading.BoundedSemaphore(1 + _PREPARED_AHEAD)
+        self._prepare_gate = threading.BoundedSemaphore(_PREPARED_AHEAD)
         # Re-entrant: a fit retrying alone holds it while the backend's
         # measurement takes it again (see _cuda.QuiescedMeasurement).
         self._predict_lock = threading.RLock()
@@ -321,16 +323,23 @@ def fit(
         # The fit's CPU half (backends that split it: TabICL's preprocessing,
         # ~1/3 of a fit) runs before taking the device, overlapping the GPU
         # half of whatever fit holds it now.
-        # At most _PREPARED_AHEAD fits hold prepared (host-resident) data
-        # while waiting for the device, so queued fits don't each keep
-        # their preprocessed copies.
+        # At most _PREPARED_AHEAD large fits hold prepared (host-resident)
+        # data while waiting for the device, so queued large tables don't
+        # each keep their preprocessed copies. The permit is returned once
+        # the fit has the device, and small fits skip the gate, so neither
+        # a running fit nor large fits waiting for host memory hold up a
+        # small fit on this device.
         # Lock order everywhere: prepare gate, host slot, _fit_lock,
         # _predict_lock. Host memory before the device, so a large fit
         # waiting for host memory doesn't hold this device meanwhile.
-        with self._prepare_gate:
+        with contextlib.ExitStack() as gate:
+            if n_train * n_features >= GATED_FIT_CELLS:
+                self._prepare_gate.acquire()
+                gate.callback(self._prepare_gate.release)
             prepare = getattr(chosen, "prepare_fit", None)
             prepared = prepare(X, y, task) if callable(prepare) else None
             with self._host_slot(chosen, n_train, n_features), self._fit_lock:
+                gate.close()  # the device is ours: let the next fit prepare
                 self._make_room(estimator, estimated)
                 with self._measuring():
                     payload = self._fit_backend(
@@ -503,7 +512,10 @@ def _host_slot(self, backend: TabularICLBackend, n_rows: int, n_features: int):
         )
         estimate = getattr(backend, "host_bytes_estimate", None)
         if callable(estimate):
-            nbytes = estimate(n_rows, n_features, offloads=offloads)
+            try:
+                nbytes = estimate(n_rows, n_features, offloads=offloads)
+            except TypeError:  # an estimate from before `offloads` existed
+                nbytes = estimate(n_rows, n_features)
         else:
             nbytes = guard.large_fit_bytes if offloads else 0
         return guard.reserve(nbytes, offloads=offloads)
diff --git a/src/tabctx/memory/host.py b/src/tabctx/memory/host.py
index e3be77b..1ca3bf0 100644
--- a/src/tabctx/memory/host.py
+++ b/src/tabctx/memory/host.py
@@ -32,10 +32,14 @@
 Robustness. Reservations are kept per owner (a worker process, or 0
 in-process), so when a worker dies holding bytes -- the host OOM kill
 this guard exists to prevent -- its parent returns them (release_owner)
-instead of leaking them until the replica restarts. The shared lock is
-taken with a timeout: a worker killed while holding it (a window of
-microseconds) makes calls fail with a retryable HostMemoryUnavailableError
-and /readyz report unhealthy, never hang.
+instead of leaking them until the replica restarts. Across processes the
+shared lock is a plain lock that waiters poll (a process killed while
+waiting on an mp.Condition would block every later notify forever), and
+every acquire has a timeout: a worker killed while holding it (a window
+of microseconds) makes calls fail with a retryable
+HostMemoryUnavailableError and /readyz report unhealthy, never hang. A
+call's release never raises: its result stands, and a reservation that
+can't be returned is logged.
 """
 
 from __future__ import annotations
@@ -43,6 +47,7 @@
 import contextlib
 import logging
 import threading
+import time
 from collections.abc import Iterator
 from pathlib import Path
 from typing import Any
@@ -59,12 +64,14 @@
 # reserve (8 GiB) for the rest of serving.
 MIN_GUARDED_BYTES = 1 * GiB
 LOCK_TIMEOUT_S = 10.0
-# Waiters re-check at least this often (a dead owner's release, a timeout).
+# In-process waiters are notified, and re-check at least this often.
 _POLL_S = 1.0
+# Across processes waiters poll (see HostMemoryGuard.__init__).
+_MP_POLL_S = 0.02
 # Per owner: reserved bytes, active calls, offloading calls, then queue
-# slots for waiting calls (a worker serves at most 16 calls at once).
-_QUEUE_SLOTS = 32
-_OWNER_BLOCK = 3 + _QUEUE_SLOTS
+# slots for its waiting calls (a worker serves at most 16 calls at once;
+# an in-process guard shared by several devices' engines needs more).
+QUEUE_SLOTS = 32
 
 
 CGROUP_LIMIT_PATHS = (
@@ -137,6 +144,7 @@ def __init__(
         owners: int = 1,
         min_guarded_bytes: int = MIN_GUARDED_BYTES,
         lock_timeout_s: float = LOCK_TIMEOUT_S,
+        queue_slots: int = QUEUE_SLOTS,
     ) -> None:
         self.budget_bytes = budget_bytes
         self.large_fit_bytes = large_fit_bytes
@@ -154,12 +162,20 @@ def __init__(
         self._lock_timeout = lock_timeout_s
         self._owners = max(1, owners)
         self._owner = 0
-        size = 1 + self._owners * _OWNER_BLOCK
+        self._block_size = 3 + max(1, queue_slots)
+        size = 1 + self._owners * self._block_size
         if mp_context is None:
+            # In-process: a condition variable (no process can die holding it).
             self._cond: Any = threading.Condition()
+            self._lock: Any = self._cond
             self._state: Any = [0] * size
         else:
-            self._cond = mp_context.Condition(mp_context.Lock())
+            # Across processes: a plain lock, and waiters poll. mp.Condition
+            # is unsafe here: a process killed while waiting leaves its
+            # sleeper count behind, and the next notify_all blocks forever
+            # (with the lock held) waiting for it to wake.
+            self._cond = None
+            self._lock = mp_context.Lock()
             self._state = mp_context.Array("q", size, lock=False)
 
     def bind(self, owner: int) -> None:
@@ -184,40 +200,39 @@ def reserve(self, nbytes: int | None, offloads: bool = False) -> Iterator[None]:
             yield
             return
         state, base = self._state, self._block(self._owner)
-        with self._locked():
+        with self._locked() as held:
             ticket = state[0] = state[0] + 1
             slot = None
             try:
                 while slot is None:
                     slot = self._park(base, ticket)
                     if slot is None:
-                        self._cond.wait(_POLL_S)
+                        self._wait(held)
                 while not (self._head() == ticket and self._admits(nbytes, offloads)):
-                    self._cond.wait(_POLL_S)
+                    self._wait(held)
             finally:
                 if slot is not None:
+                    # Unpark even without the lock (a lost reacquire): an
+                    # orphaned ticket would stay the head of the queue.
                     state[slot] = 0
-                self._cond.notify_all()  # the next waiter is now the head
+                    if held[0]:
+                        self._notify()  # the next waiter is now the head
             state[base] += nbytes
             state[base + 1] += 1
             state[base + 2] += int(offloads)
         try:
             yield
         finally:
-            with self._locked():
-                state[base] = max(0, state[base] - nbytes)
-                state[base + 1] = max(0, state[base + 1] - 1)
-                state[base + 2] = max(0, state[base + 2] - int(offloads))
-                self._cond.notify_all()
+            self._release(base, nbytes, offloads)
 
     def release_owner(self, owner: int) -> None:
         """Return everything a dead owner held or queued for."""
         base = self._block(owner)
         try:
             with self._locked():
-                for i in range(base, base + _OWNER_BLOCK):
+                for i in range(base, base + self._block_size):
                     self._state[i] = 0
-                self._cond.notify_all()
+                self._notify()
         except HostMemoryUnavailableError:
             logger.error("could not release host reservations of worker %d", owner)
 
@@ -233,7 +248,7 @@ def stats(self) -> dict[str, int]:
                 "calls_waiting": sum(
                     1
                     for b in blocks
-                    for i in range(b + 3, b + _OWNER_BLOCK)
+                    for i in range(b + 3, b + self._block_size)
                     if self._state[i]
                 ),
                 "large_fit_slots": self.slots,
@@ -241,24 +256,59 @@ def stats(self) -> dict[str, int]:
 
     # ---- internals (callers hold the lock) ----------------------------
 
+    def _release(self, base: int, nbytes: int, offloads: bool) -> None:
+        """Return a call's reservation. Never raises: the call itself has
+        succeeded (a fit's context may already be cached), so a stuck lock
+        is logged -- and seen by /readyz, whose stats() fails -- rather
+        than turning that result into an error."""
+        state = self._state
+        try:
+            with self._locked():
+                state[base] = max(0, state[base] - nbytes)
+                state[base + 1] = max(0, state[base + 1] - 1)
+                state[base + 2] = max(0, state[base + 2] - int(offloads))
+                self._notify()
+        except HostMemoryUnavailableError:
+            logger.error("host memory reservation of %d bytes not returned", nbytes)
+
     @contextlib.contextmanager
-    def _locked(self) -> Iterator[None]:
-        if not self._cond.acquire(True, self._lock_timeout):
-            raise HostMemoryUnavailableError(
-                f"host memory accounting lock not acquired in {self._lock_timeout}s "
-                "(a worker may have died holding it); retry"
-            )
+    def _locked(self) -> Iterator[list[bool]]:
+        """Hold the lock; yields a flag _wait clears if it can't get the
+        lock back (then there is nothing to release)."""
+        if not self._lock.acquire(True, self._lock_timeout):
+            raise self._stuck()
+        held = [True]
         try:
-            yield
+            yield held
         finally:
-            self._cond.release()
+            if held[0]:
+                self._lock.release()
+
+    def _wait(self, held: list[bool]) -> None:
+        if self._cond is not None:
+            self._cond.wait(_POLL_S)
+            return
+        self._lock.release()
+        time.sleep(_MP_POLL_S)
+        if not self._lock.acquire(True, self._lock_timeout):
+            held[0] = False
+            raise self._stuck()
+
+    def _notify(self) -> None:
+        if self._cond is not None:
+            self._cond.notify_all()
+
+    def _stuck(self) -> HostMemoryUnavailableError:
+        return HostMemoryUnavailableError(
+            f"host memory accounting lock not acquired in {self._lock_timeout}s "
+            "(a worker may have died holding it); retry"
+        )
 
-    @staticmethod
-    def _block(owner: int) -> int:
-        return 1 + owner * _OWNER_BLOCK
+    def _block(self, owner: int) -> int:
+        return 1 + owner * self._block_size
 
     def _park(self, base: int, ticket: int) -> int | None:
-        for i in range(base + 3, base + _OWNER_BLOCK):
+        for i in range(base + 3, base + self._block_size):
             if self._state[i] == 0:
                 self._state[i] = ticket
                 return i
@@ -269,7 +319,7 @@ def _head(self) -> int:
             (
                 self._state[i]
                 for o in range(self._owners)
-                for i in range(self._block(o) + 3, self._block(o) + _OWNER_BLOCK)
+                for i in range(self._block(o) + 3, self._block(o) + self._block_size)
                 if self._state[i]
             ),
             default=0,
diff --git a/src/tabctx/serve/factory.py b/src/tabctx/serve/factory.py
index 450a92d..b9b2690 100644
--- a/src/tabctx/serve/factory.py
+++ b/src/tabctx/serve/factory.py
@@ -92,7 +92,11 @@
     DEFAULT_HARD_CEILING_BYTES,
     MemoryEstimator,
 )
-from tabctx.memory.host import HostMemoryGuard, detect_host_budget_bytes
+from tabctx.memory.host import (
+    QUEUE_SLOTS,
+    HostMemoryGuard,
+    detect_host_budget_bytes,
+)
 from tabctx.memory.requirement import ResidentRateModel, WorkingFloorModel
 from tabctx.multidevice import MultiDeviceEngine
 from tabctx.serve.uploads import DEFAULT_MAX_UPLOAD_BYTES, DEFAULT_TTL_S
@@ -669,8 +673,13 @@ def build_engine(settings: ServeSettings | None = None) -> BuiltEngine:
     host_budget = None if fake_only else detect_host_budget_bytes()
     if multi and settings.worker_processes:
         return _build_process_engine(settings, devices, gpu_capacity, host_budget)
-    # One host-memory guard per process: host RAM is shared by every device.
-    host_guard = HostMemoryGuard(host_budget) if host_budget else None
+    # One host-memory guard per process (host RAM is shared by every
+    # device), with queue slots for each device's engine's waiting calls.
+    host_guard = (
+        HostMemoryGuard(host_budget, queue_slots=QUEUE_SLOTS * len(devices))
+        if host_budget
+        else None
+    )
     estimators: dict[str, MemoryEstimator] = {}
     shards = [
         build_shard(
diff --git a/tests/unit/test_engine_scheduling.py b/tests/unit/test_engine_scheduling.py
index a1803aa..779f714 100644
--- a/tests/unit/test_engine_scheduling.py
+++ b/tests/unit/test_engine_scheduling.py
@@ -351,10 +351,13 @@ def test_an_infeasible_fit_is_rejected_before_preparing():
     assert backend.prepared_during == []
 
 
-def test_only_one_fit_prepares_ahead_of_the_device():
+def test_only_one_large_fit_prepares_ahead_of_the_device(monkeypatch):
     """Prepared data lives in host memory until its fit gets the device:
-    queued fits beyond one ahead wait before preparing, so N queued large
-    tables don't hold N preprocessed copies at once."""
+    queued large fits beyond one ahead wait before preparing, so N queued
+    large tables don't hold N preprocessed copies at once."""
+    import tabctx.engine as engine_module
+
+    monkeypatch.setattr(engine_module, "GATED_FIT_CELLS", 10)
     backend = PreparingFake()
     engine = TabctxEngine(
         backend=backend, estimator=CellEstimator(), cache=ContextCacheManager(CAPACITY)
@@ -371,3 +374,22 @@ def test_only_one_fit_prepares_ahead_of_the_device():
         t.join(5)
     assert len(backend.prepared_during) == 4
     assert all(engine._cache.holds(f"d{i}") for i in range(4))
+
+
+def test_small_fits_skip_the_prepare_gate(monkeypatch):
+    """A large fit holding the gate (e.g. waiting for host memory) must not
+    hold up a small fit's preparation on the same device."""
+    import tabctx.engine as engine_module
+
+    monkeypatch.setattr(engine_module, "GATED_FIT_CELLS", 50)
+    backend = PreparingFake()
+    engine = TabctxEngine(
+        backend=backend, estimator=CellEstimator(), cache=ContextCacheManager(CAPACITY)
+    )
+    backend.engine = engine
+    engine._prepare_gate.acquire()  # a large fit holds the only permit
+    try:
+        engine.fit(*table(10), dataset_id="small")  # 10 cells: ungated
+        assert engine._cache.holds("small")
+    finally:
+        engine._prepare_gate.release()
diff --git a/tests/unit/test_host_memory.py b/tests/unit/test_host_memory.py
index 9c02b2f..d8a6436 100644
--- a/tests/unit/test_host_memory.py
+++ b/tests/unit/test_host_memory.py
@@ -251,6 +251,62 @@ def parent_call():
             child.kill()
 
 
+def _wait_in_child(guard, owner, nbytes, queued):
+    guard.bind(owner)
+    queued.set()
+    with guard.reserve(nbytes):
+        pass
+
+
+def test_a_process_killed_while_waiting_does_not_wedge_the_guard():
+    """With mp.Condition, a process killed in wait() left its sleeper count
+    behind and the next notify_all blocked forever, holding the lock."""
+    import multiprocessing as mp
+
+    ctx = mp.get_context("spawn")
+    guard = HostMemoryGuard(10 * GiB, mp_context=ctx, owners=2, lock_timeout_s=5)
+    queued = ctx.Event()
+    done = threading.Event()
+    with guard.reserve(8 * GiB):  # blocks the child's 4 GiB
+        child = ctx.Process(
+            target=_wait_in_child, args=(guard, 1, 4 * GiB, queued), daemon=True
+        )
+        child.start()
+        assert queued.wait(60)
+        _wait_for(lambda: guard.stats()["calls_waiting"] == 1, timeout=30)
+        child.kill()
+        child.join(10)
+
+        def after():
+            guard.release_owner(1)
+            done.set()
+
+        threading.Thread(target=after, daemon=True).start()
+        assert done.wait(10), "release_owner hung after a waiter was killed"
+        assert guard.stats()["calls_waiting"] == 0
+    with guard.reserve(9 * GiB):  # and the guard still admits
+        assert guard.stats()["reserved_bytes"] == 9 * GiB
+
+
+def test_a_stuck_lock_at_release_does_not_fail_the_call(caplog):
+    guard = HostMemoryGuard(10 * GiB, large_fit_bytes=GiB, lock_timeout_s=0.2)
+    holding, done = threading.Event(), threading.Event()
+
+    def hold_lock():
+        with guard._cond:
+            holding.set()
+            done.wait(5)
+
+    with caplog.at_level("ERROR"), guard.reserve(2 * GiB):  # the call succeeds...
+        t = threading.Thread(target=hold_lock, daemon=True)
+        t.start()
+        assert holding.wait(5)
+        # ...and its release, unable to get the lock, logs instead of raising
+    assert "not returned" in caplog.text
+    done.set()
+    t.join(5)
+
+
 def test_bind_rejects_unknown_owners():
     with pytest.raises(ValueError):
         HostMemoryGuard(GiB, owners=2).bind(2)
@@ -434,3 +490,17 @@ def host_bytes_estimate(self, n_rows, n_features, *, offloads=False):
     seen = _spy(guard)
     engine.fit([[1.0, 2.0]] * 3, ["a"] * 3, dataset_id="d")
     assert seen == [(84, True)]
+
+
+def test_a_two_argument_host_estimate_still_works():
+    """Backends written before the engine passed `offloads=`."""
+
+    class OldStyle(FakeBackend):
+        def host_bytes_estimate(self, n_rows, n_features):
+            return 5 * n_rows * n_features
+
+    guard = HostMemoryGuard(GiB, large_fit_bytes=GiB, min_guarded_bytes=0)
+    engine = _engine(OldStyle(bytes_hint=10), guard)
+    seen = _spy(guard)
+    engine.fit([[1.0, 2.0]] * 3, ["a"] * 3, dataset_id="d")
+    assert seen == [(30, False)]  # 6 cells: under the large-fit threshold
diff --git a/tests/unit/test_tabicl_thread_local_cache.py b/tests/unit/test_tabicl_thread_local_cache.py
index 074348e..91a8a80 100644
--- a/tests/unit/test_tabicl_thread_local_cache.py
+++ b/tests/unit/test_tabicl_thread_local_cache.py
@@ -112,6 +112,7 @@ def test_compiled_blocks_install_once():
         assert install_compiled_blocks()
         compiled = layers.MultiheadAttentionBlock.forward
         assert compiled is not original
+        assert compiled._tabctx_eager is original
         assert install_compiled_blocks()
         assert layers.MultiheadAttentionBlock.forward is compiled  # idempotent
     finally:

From e03917131920c774624896aa379c075ee75edcb7 Mon Sep 17 00:00:00 2001
From: Amrit Krishnan <amrit110@gmail.com>
Date: Sun, 18 Oct 2026 00:46:02 -0400
Subject: [PATCH 31/39] feat(multidevice): pluggable placement policy and
 residency events (#39)

The router keeps the mechanics (snapshots, per-id locks, re-fit order)
and asks a PlacementPolicy only the decision, from a lock-free DeviceView
per device (headroom, the fit's need there, activity, liveness). The
previous rule is the default LeastActiveThatFits, which now also skips
dead workers explicitly. Placements and removals are ResidencyEvents fanned
out to sinks (a failing sink is logged, never fails the request);
/v1/tabctx/residency serves the latest. A platform co-scheduling LLMs and
tabular models (Ray Serve now, Dynamo next) can plug in its own policy
and follow placements without polling.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 CHANGELOG.md                        |   1 +
 README.md                           |   6 ++
 ROADMAP.md                          |  10 ++-
 src/tabctx/events.py                |  82 +++++++++++++++++++++
 src/tabctx/multidevice.py           |  98 +++++++++++++++----------
 src/tabctx/placement.py             |  78 ++++++++++++++++++++
 src/tabctx/serve/app.py             |  15 ++++
 tests/unit/test_placement_events.py | 108 ++++++++++++++++++++++++++++
 8 files changed, 357 insertions(+), 41 deletions(-)
 create mode 100644 src/tabctx/events.py
 create mode 100644 src/tabctx/placement.py
 create mode 100644 tests/unit/test_placement_events.py

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7ace6c5..c6930b5 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,6 +8,7 @@
 - **One worker process per GPU** (#39, `TABCTX_WORKER_PROCESSES`, default on for multi-GPU replicas). Fits are host-bound, so threads of one process driving several GPUs contend for the GIL: on 4x A100, four concurrent fits ran at 2.1-3.1x aggregate from threads and 3.96-4.3x from processes (`benchmarks/bench_scaling.py`). `serve/workers.py` runs each device's engine in its own process behind a `ProcessShard` proxy; the router reads versioned residency/headroom snapshots that every reply carries, so its decisions stay lock-free. A dead worker fails its calls, is skipped by placement, and fails `/readyz` and Ray Serve's `check_health` so the replica restarts. Nothing depends on Ray. Under Ray Serve on 4x A100 with 6 mixed users: 2.58 -> 5.0 predicts/s, small-fit p50 1.6 -> 0.6 s.
 - **Light predicts run beside a fit** on the same GPU (#39). A predict against a resident context with at most 2,048 rows no longer waits for a fit on its device (it waited up to 266 s on the L4 rehearsal): ~0.1 s p50 during a 40k x 100 fit on an A100, at ~15% fit slowdown. Restores and large predicts still wait; contexts in use are pinned against evict-ahead (and evicted with predicts held off when nothing else makes room); a fit that runs out of memory beside a predict retries alone; allocator measurements are taken with predicts quiesced, and memory freed during a fit is released only after it has measured itself.
 - **TabPFN-3.5 requirement calibration** (A100, 20 shapes up to 200k x 10 and 30k x 1,000). TabPFN also sizes its working memory from what is free (20k x 100 peaks at 7.3 GiB unconstrained but runs within 2x its time on 1.1 GiB), so admission bounds resident + working floor for it too, instead of TabICL's peak formula: up to 100k rows at 100 features and 80k at 1,000 on an A40 (was 5k and 500 on an L4). Its per-row resident rate grows with features, which `ResidentRateModel(width_dependent=True)` respects.
+- **Pluggable placement and residency events** (#39). The multi-device router asks a `PlacementPolicy` (`placement.py`) which device takes a fit, from a lock-free view of each device (headroom, the fit's estimated need there, activity, liveness); the previous rule is the default `LeastActiveThatFits`, which now also skips dead workers. The router emits `ResidencyEvent`s (`events.py`) when it places a context and when it removes a copy, to any sink; `/v1/tabctx/residency` serves the latest. A platform that co-schedules LLMs and tabular models can supply its own policy and follow placements without polling.
 - `benchmarks/bench_scaling.py` (threads vs processes across GPUs), `bench_overlap.py` (predict latency and correctness during a fit), `gke/serve_workload.sh` (Ray Serve + workload inside the devbox).
 - `benchmarks/bench_workload.py`: concurrent users looping fit-then-predict over a small/medium/large table mix (large ones via CSV upload), with per-class latency percentiles, 413s, errors and the per-device spread.
 - `benchmarks/sweep_admission.py`: the table-size sweep (rows up to the first 413 per feature count, fit and predict timed).
diff --git a/README.md b/README.md
index f989823..54a8dd3 100644
--- a/README.md
+++ b/README.md
@@ -233,6 +233,12 @@ behind the same API. No client changes and no routing header needed.
   against a resident context runs beside a fit on the same GPU (~0.1 s
   instead of the fit's full duration); larger predicts and restores of
   spilled contexts wait their turn.
+- **Placement is pluggable.** The rule above is the default
+  `PlacementPolicy` (`tabctx.placement`); pass your own to
+  `MultiDeviceEngine(placement=...)` to steer tabular work, e.g. away
+  from a GPU a platform scheduler wants for an LLM. Each placement and
+  removal is a residency event (`tabctx.events`) for any subscriber;
+  `GET /v1/tabctx/residency` returns the latest.
 - **The replica must be given the GPUs.** Under Ray Serve that means
   `ray_actor_options.num_gpus: N`, or, for a `num_gpus: 0` actor on a
   GPU pod (as in inference-platform's prod config),
diff --git a/ROADMAP.md b/ROADMAP.md
index 578fe47..1eec921 100644
--- a/ROADMAP.md
+++ b/ROADMAP.md
@@ -112,9 +112,13 @@ no FP8).
     the graph), so it stays off until those are traced down.
 14. **Serving-framework-agnostic core** ([#39](https://github.com/VectorInstitute/tabctx/issues/39)): the router, per-GPU workers
     and scheduling are plain Python with no Ray dependency; a Dynamo
-    worker hosts the same core. Next: residency events and pluggable
-    placement so a platform scheduler can co-optimize LLMs and tabular
-    models.
+    worker hosts the same core. Placement is a pluggable policy
+    (`placement.py`: the router hands it a lock-free view of each device)
+    and the router emits residency events (`events.py`: placed/removed,
+    to any sink; `/v1/tabctx/residency` serves the latest), so a platform
+    scheduler can follow where contexts live and steer tabular work
+    around GPUs it wants for LLMs. Next: a Dynamo worker and event-plane
+    sink once the platform moves.
 
 PyTorch-level first: tiling and chunking, `torch.compile` (generated Triton
 kernels, CUDA graphs via `reduce-overhead`), FlexAttention for paged and
diff --git a/src/tabctx/events.py b/src/tabctx/events.py
new file mode 100644
index 0000000..610b103
--- /dev/null
+++ b/src/tabctx/events.py
@@ -0,0 +1,82 @@
+"""Residency events: where contexts live, pushed to whoever schedules (#39).
+
+A platform that co-schedules LLMs and tabular models needs to know which
+GPU holds which tenant's context without polling every replica. The
+multi-device router emits a ``ResidencyEvent`` when a fit places a context
+on a device and when a copy is removed (re-fit elsewhere, explicit evict).
+Sinks are plain callables, so any transport plugs in: a log, a metrics
+counter, a Ray actor, a Dynamo event plane publisher.
+
+Events report the router's decisions, not every cache-internal move:
+capacity eviction to the spill tier and restores stay on their device,
+so a context's device doesn't change until the router says so.
+
+Sinks run on the request's thread after the change, so they must be
+fast and must not raise (one that does is logged and skipped).
+"""
+
+from __future__ import annotations
+
+import logging
+import threading
+import time
+from collections import deque
+from collections.abc import Callable
+from dataclasses import asdict, dataclass, field
+from typing import Literal
+
+logger = logging.getLogger(__name__)
+
+EventKind = Literal["placed", "removed"]
+
+
+@dataclass(frozen=True)
+class ResidencyEvent:
+    kind: EventKind
+    dataset_id: str
+    device: str
+    backend: str | None = None
+    timestamp: float = field(default_factory=time.time)
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+
+EventSink = Callable[[ResidencyEvent], None]
+
+
+class EventBus:
+    """Fans events out to sinks; a failing sink never fails the request."""
+
+    def __init__(self, sinks: list[EventSink] | None = None) -> None:
+        self._sinks: list[EventSink] = list(sinks or [])
+        self._lock = threading.Lock()
+
+    def subscribe(self, sink: EventSink) -> None:
+        with self._lock:
+            self._sinks.append(sink)
+
+    def emit(self, event: ResidencyEvent) -> None:
+        with self._lock:
+            sinks = list(self._sinks)
+        for sink in sinks:
+            try:
+                sink(event)
+            except Exception:  # noqa: BLE001 -- observers must not break serving
+                logger.exception("residency event sink failed on %s", event)
+
+
+class RecentEvents:
+    """A bounded in-memory log of the latest events (a sink)."""
+
+    def __init__(self, maxlen: int = 256) -> None:
+        self._events: deque[ResidencyEvent] = deque(maxlen=maxlen)
+        self._lock = threading.Lock()
+
+    def __call__(self, event: ResidencyEvent) -> None:
+        with self._lock:
+            self._events.append(event)
+
+    def snapshot(self) -> list[dict]:
+        with self._lock:
+            return [e.to_dict() for e in self._events]
diff --git a/src/tabctx/multidevice.py b/src/tabctx/multidevice.py
index 9b1f382..803a454 100644
--- a/src/tabctx/multidevice.py
+++ b/src/tabctx/multidevice.py
@@ -7,22 +7,13 @@
 to that device, its own cache, lock and spill tier -- and this class
 routes between them:
 
-- fit() places each context with a rule aimed at the usual workload --
-  a few concurrent users, each fitting a table (small or large) and then
-  predicting against it repeatedly:
-    1. Devices where the table's estimated peak fits WITHOUT evicting
-       anything are preferred; among them, the least active one wins
-       (fewest calls in flight, then fewest started recently), ties to
-       the most headroom. Each device runs one fit at a time, so
-       spreading active users across devices is what lets them run
-       concurrently.
-    2. If no device can take it without eviction (a large table, or a
-       full replica), it goes to the device with the most headroom,
-       evicting the fewest contexts.
-  Each device's own usage-aware admission gate then has the final say,
-  exactly as on a single GPU. A table's ceiling is therefore one
-  device's memory; what more devices buy is N x the resident cache and
-  concurrency.
+- fit() places each context where a PlacementPolicy says (placement.py;
+  the default, LeastActiveThatFits, prefers devices where the table fits
+  without eviction, the least active among them, else the one with most
+  headroom). Each device's own usage-aware admission gate then has the
+  final say, exactly as on a single GPU. A table's ceiling is therefore
+  one device's memory; what more devices buy is N x the resident cache
+  and concurrency.
 - predict() runs on the device that holds the context, found from each
   device's residency and spill index (so a context spilled before a
   restart is found on, and restored to, its own device).
@@ -43,6 +34,9 @@
 single-GPU fit would. Serializing placement would cost the concurrency
 this class exists for.
 
+Residency events (events.py) report each placement and removal, so a
+platform scheduler can follow where contexts live without polling.
+
 Same public surface as TabctxEngine, so the serving layer and batching
 use either interchangeably.
 """
@@ -58,6 +52,13 @@
 
 from tabctx.engine import _validate_fit_input
 from tabctx.errors import DatasetNotFoundError
+from tabctx.events import EventBus, EventKind, ResidencyEvent
+from tabctx.placement import (
+    DeviceView,
+    FitRequest,
+    LeastActiveThatFits,
+    PlacementPolicy,
+)
 from tabctx.types import ArrayLike, EngineStats, PredictOutcome, Task
 
 # How long a started call counts toward a device's recent activity.
@@ -146,6 +147,8 @@ def __init__(
         devices: list[str],
         activity_window_s: float = ACTIVITY_WINDOW_S,
         clock: Callable[[], float] = time.monotonic,
+        placement: PlacementPolicy | None = None,
+        events: EventBus | None = None,
     ) -> None:
         if len(shards) < 2 or len(shards) != len(devices):
             raise ValueError("need one engine per device, and at least two devices")
@@ -161,6 +164,8 @@ def __init__(
         self._inflight = [0] * len(shards)
         self._recent: list[deque[float]] = [deque() for _ in shards]
         self._activity_lock = threading.Lock()
+        self._placement: PlacementPolicy = placement or LeastActiveThatFits()
+        self.events = events or EventBus()
 
     @property
     def devices(self) -> list[str]:
@@ -210,19 +215,30 @@ def _activity(self) -> list[tuple[int, int]]:
 
     # ---- placement and lookup (lock-free snapshots) ----------------
 
-    def _place(self, n_train: int, n_features: int, backend: str) -> int:
-        headrooms = self.device_headroom_bytes(backend)
-        # Each shard prices the fit with its own estimator (worker
-        # processes learn separately; in-process shards share one).
-        needs = [
-            s.estimate_fit_bytes(n_train, n_features, backend) for s in self._shards
-        ]
+    def _place(self, request: FitRequest) -> int:
+        """Ask the policy, from lock-free snapshots of every device."""
+        headrooms = self.device_headroom_bytes(request.backend)
         activity = self._activity()
-        indices = range(len(self._shards))
-        fits = [i for i in indices if headrooms[i] >= needs[i]]
-        if fits:
-            return min(fits, key=lambda i: (activity[i], -headrooms[i], i))
-        return max(indices, key=lambda i: (headrooms[i] - needs[i], -i))
+        views = [
+            DeviceView(
+                index=i,
+                device=self._devices[i],
+                headroom_bytes=headrooms[i],
+                # Each shard prices the fit with its own estimator (worker
+                # processes learn separately; in-process shards share one).
+                need_bytes=s.estimate_fit_bytes(
+                    request.n_train, request.n_features, request.backend
+                ),
+                in_flight=activity[i][0],
+                recent_calls=activity[i][1],
+                alive=getattr(s, "alive", True),
+            )
+            for i, s in enumerate(self._shards)
+        ]
+        index = self._placement.choose(request, views)
+        if not 0 <= index < len(self._shards):
+            raise ValueError(f"placement policy chose device {index} of {len(views)}")
+        return index
 
     def _locate(self, dataset_id: str) -> list[int]:
         """Devices holding the id: resident first, then spilled copies."""
@@ -247,26 +263,31 @@ def fit(
     ) -> str:
         n_features = _validate_fit_input(X, y)
         name = self._shards[0].resolve_backend_name(backend)
+        request = FitRequest(len(X), n_features, name, dataset_id)
         if dataset_id is None:
-            return self._fit_on(
-                self._place(len(X), n_features, name),
-                X,
-                y,
-                task,
-                None,
-                name,
-                feature_names,
-            )
+            index = self._place(request)
+            resolved = self._fit_on(index, X, y, task, None, name, feature_names)
+            self._emit("placed", resolved, index, name)
+            return resolved
         with self._id_locks.hold(dataset_id):
-            index = self._place(len(X), n_features, name)
+            index = self._place(request)
             resolved = self._fit_on(index, X, y, task, dataset_id, name, feature_names)
+            self._emit("placed", resolved, index, name)
             # Only now that the new context is in place: drop copies on
             # other devices (a failed fit above leaves the old one serving).
             for i in self._locate(dataset_id):
                 if i != index:
                     self._shards[i].evict(dataset_id)
+                    self._emit("removed", dataset_id, i)
             return resolved
 
+    def _emit(
+        self, kind: EventKind, dataset_id: str, index: int, backend: str | None = None
+    ) -> None:
+        self.events.emit(
+            ResidencyEvent(kind, dataset_id, self._devices[index], backend)
+        )
+
     def _fit_on(self, index, X, y, task, dataset_id, backend, feature_names) -> str:
         with self._active(index):
             return self._shards[index].fit(
@@ -325,6 +346,7 @@ def evict(self, dataset_id: str) -> None:
         with self._id_locks.hold(dataset_id):
             for index in self._locate(dataset_id):
                 self._shards[index].evict(dataset_id)
+                self._emit("removed", dataset_id, index)
 
     def stats(self) -> EngineStats:
         per_device = self.device_stats()
diff --git a/src/tabctx/placement.py b/src/tabctx/placement.py
new file mode 100644
index 0000000..3feff6a
--- /dev/null
+++ b/src/tabctx/placement.py
@@ -0,0 +1,78 @@
+"""Where a fit goes among a replica's devices: a pluggable policy (#39).
+
+The router (multidevice.py) owns the mechanics -- snapshots, per-id
+locks, re-fit ordering -- and asks a ``PlacementPolicy`` only the
+decision: given each device's view, which one takes this fit. Keeping
+the decision separate lets a platform that schedules LLMs and tabular
+models together (Ray Serve today, Dynamo tomorrow) supply its own policy,
+e.g. one that avoids a GPU it is about to give to an LLM, without
+touching the router.
+
+A view is a lock-free snapshot, so a policy never waits on a device's
+work, and it must be cheap: it runs on every fit.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import Protocol
+
+
+@dataclass(frozen=True)
+class FitRequest:
+    n_train: int
+    n_features: int
+    backend: str
+    dataset_id: str | None = None
+
+
+@dataclass(frozen=True)
+class DeviceView:
+    """One device as a policy sees it."""
+
+    index: int
+    device: str
+    headroom_bytes: int  # what its admission gate offers now
+    need_bytes: int  # this fit's estimated peak there
+    in_flight: int  # calls running on it
+    recent_calls: int  # calls started within the activity window
+    alive: bool = True
+
+    @property
+    def fits_without_eviction(self) -> bool:
+        return self.alive and self.headroom_bytes >= self.need_bytes
+
+
+class PlacementPolicy(Protocol):
+    def choose(self, request: FitRequest, devices: Sequence[DeviceView]) -> int:
+        """The index of the device that takes the fit."""
+        ...
+
+
+class LeastActiveThatFits:
+    """The default, aimed at a few concurrent users fitting small and large
+    tables and predicting against them:
+
+    1. Devices where the fit's peak fits WITHOUT evicting anything win;
+       among them the least active (fewest calls in flight, then fewest
+       started recently), ties to the most headroom. Each device runs one
+       fit at a time, so spreading active users is what lets them run
+       concurrently.
+    2. Otherwise (a large table, or a full replica) the device with the
+       most headroom left after the fit, i.e. the fewest evictions.
+
+    Dead devices are never chosen while a live one exists."""
+
+    def choose(self, request: FitRequest, devices: Sequence[DeviceView]) -> int:
+        del request
+        live = [d for d in devices if d.alive] or list(devices)
+        fits = [d for d in live if d.fits_without_eviction]
+        if fits:
+            best = min(
+                fits,
+                key=lambda d: (d.in_flight, d.recent_calls, -d.headroom_bytes, d.index),
+            )
+        else:
+            best = max(live, key=lambda d: (d.headroom_bytes - d.need_bytes, -d.index))
+        return best.index
diff --git a/src/tabctx/serve/app.py b/src/tabctx/serve/app.py
index 4529729..59eef16 100644
--- a/src/tabctx/serve/app.py
+++ b/src/tabctx/serve/app.py
@@ -45,6 +45,7 @@
     UploadTooLargeError,
 )
 from tabctx.batching import CoalescingPredictor
+from tabctx.events import RecentEvents
 from tabctx.multidevice import MultiDeviceEngine
 from tabctx.serve.affinity import resolve_dataset_id, session_id_from_headers
 from tabctx.serve.csv_io import parse_features_csv, parse_train_csv
@@ -226,6 +227,12 @@ def __init__(self) -> None:
         self._devices = built.devices
         self._device = built.device
         self._gpu_capacity_bytes = built.gpu_capacity_bytes
+        # Residency events from the multi-device router (events.py): the
+        # latest, for /v1/tabctx/residency; a platform adds its own sinks.
+        self._residency = RecentEvents()
+        bus = getattr(self._engine, "events", None)
+        if bus is not None:
+            bus.subscribe(self._residency)
         # Same-context predict coalescing (see batching.py): concurrent
         # requests against one cached context share a single backend call.
         self._predictor = CoalescingPredictor(
@@ -524,6 +531,14 @@ def models(self):
             ]
         }
 
+    @fastapi_app.get("/v1/tabctx/residency")
+    def residency(self):
+        """Where contexts were placed and removed, latest last (multi-GPU
+        replicas): the router's decisions, for a platform scheduler or an
+        operator. Empty on a single-GPU replica, whose contexts all live
+        on its one device."""
+        return {"devices": list(self._devices), "events": self._residency.snapshot()}
+
     @fastapi_app.get("/v1/tabctx/limits")
     def limits(self):
         """Capability discovery (inspired by PriorLabs' get_model_limits):
diff --git a/tests/unit/test_placement_events.py b/tests/unit/test_placement_events.py
new file mode 100644
index 0000000..bd98d86
--- /dev/null
+++ b/tests/unit/test_placement_events.py
@@ -0,0 +1,108 @@
+"""Pluggable placement (placement.py) and residency events (events.py)."""
+
+import pytest
+
+from tabctx.events import EventBus, RecentEvents, ResidencyEvent
+from tabctx.multidevice import MultiDeviceEngine
+from tabctx.placement import DeviceView, FitRequest, LeastActiveThatFits
+from tests.unit.test_multidevice import build, table, where
+
+REQ = FitRequest(n_train=10, n_features=1, backend="m")
+
+
+def view(index, headroom, need=100, in_flight=0, recent=0, alive=True):
+    return DeviceView(index, f"cuda:{index}", headroom, need, in_flight, recent, alive)
+
+
+class TestDefaultPolicy:
+    def test_least_active_device_that_fits(self):
+        views = [view(0, 1000, in_flight=1), view(1, 500), view(2, 900, recent=3)]
+        assert LeastActiveThatFits().choose(REQ, views) == 1
+
+    def test_ties_go_to_most_headroom(self):
+        assert LeastActiveThatFits().choose(REQ, [view(0, 500), view(1, 900)]) == 1
+
+    def test_nothing_fits_takes_most_headroom_left(self):
+        views = [view(0, 50, need=200), view(1, 80, need=100)]
+        assert LeastActiveThatFits().choose(REQ, views) == 1
+
+    def test_dead_devices_are_skipped(self):
+        views = [view(0, 10**9, alive=False), view(1, 10, need=100)]
+        assert LeastActiveThatFits().choose(REQ, views) == 1
+
+
+class Pinned:
+    """A platform policy: keep tabular work off device 0."""
+
+    def __init__(self):
+        self.seen: list[tuple[FitRequest, list[DeviceView]]] = []
+
+    def choose(self, request, devices):
+        self.seen.append((request, list(devices)))
+        return len(devices) - 1
+
+
+def test_a_custom_policy_decides_placement():
+    engine, shards, _ = build(n_devices=3)
+    policy = Pinned()
+    router = MultiDeviceEngine(shards, engine.devices, placement=policy)
+    router.fit(*table(10), dataset_id="d")
+    assert where(shards, "d") == [2]
+    request, views = policy.seen[0]
+    assert (request.n_train, request.n_features, request.dataset_id) == (10, 1, "d")
+    assert [v.device for v in views] == ["cuda:0", "cuda:1", "cuda:2"]
+    assert all(v.need_bytes > 0 and v.headroom_bytes > 0 for v in views)
+
+
+def test_a_policy_choosing_a_bad_index_is_an_error():
+    engine, shards, _ = build(n_devices=2)
+
+    class Broken:
+        def choose(self, request, devices):
+            return 7
+
+    router = MultiDeviceEngine(shards, engine.devices, placement=Broken())
+    with pytest.raises(ValueError, match="placement policy"):
+        router.fit(*table(10), dataset_id="d")
+
+
+def test_placements_and_removals_are_reported():
+    engine, shards, _ = build(n_devices=2)
+    log = RecentEvents()
+    engine.events.subscribe(log)
+    engine.fit(*table(10), dataset_id="d")
+    first = where(shards, "d")[0]
+    # Re-fit on the other device (a policy that picks it), then evict.
+    other = 1 - first
+    engine._placement = type("P", (), {"choose": lambda self, r, d: other})()
+    engine.fit(*table(10), dataset_id="d")
+    engine.evict("d")
+    kinds = [(e["kind"], e["device"]) for e in log.snapshot()]
+    assert kinds == [
+        ("placed", f"cuda:{first}"),
+        ("placed", f"cuda:{other}"),
+        ("removed", f"cuda:{first}"),
+        ("removed", f"cuda:{other}"),
+    ]
+    assert log.snapshot()[0]["backend"] == "m"
+
+
+def test_a_failing_sink_never_fails_the_request(caplog):
+    bus = EventBus()
+    seen = []
+
+    def broken(event):
+        raise RuntimeError("sink down")
+
+    bus.subscribe(broken)
+    bus.subscribe(seen.append)
+    with caplog.at_level("ERROR"):
+        bus.emit(ResidencyEvent("placed", "d", "cuda:0"))
+    assert len(seen) == 1 and "sink failed" in caplog.text
+
+
+def test_recent_events_are_bounded():
+    log = RecentEvents(maxlen=3)
+    for i in range(5):
+        log(ResidencyEvent("placed", f"d{i}", "cuda:0"))
+    assert [e["dataset_id"] for e in log.snapshot()] == ["d2", "d3", "d4"]

From 68b716323dea6b7a55df7ec2600bf909500beb0e Mon Sep 17 00:00:00 2001
From: Amrit Krishnan <amrit110@gmail.com>
Date: Sun, 18 Oct 2026 01:03:10 -0400
Subject: [PATCH 32/39] perf(tabicl): fit Yeo-Johnson power transforms on the
 GPU, all columns at once

tabicl normalizes each ensemble member with scikit-learn's Yeo-Johnson
PowerTransformer, which fits one lambda per column with scipy's bounded
Brent search: 80%+ of a fit's CPU half (40k x 100: 3.2 of 3.9 s).
FastPowerTransformer fits every column at once with torch on the
backend's GPU -- scipy's objective, per-column overflow bounds and
tolerance; a 64-point grid then golden-section refinement -- and keeps
scikit-learn's transform scaler, so downstream (transform, spill
pickling) sees an ordinary PowerTransformer.

A100: 10k x 50 0.78 -> 0.08 s, 40k x 100 3.38 -> 0.25 s, 20k x 200 4.39 ->
0.23 s; lambdas within one float32 ulp of scikit-learn's; accuracy/AUC/R^2
unchanged on the evaluation datasets (eval_kv_precision
--fast-preprocessing); 0 wrong answers with predicts beside a fit. An
isolated 40k x 100 fit: 11.2 s eager -> 7.7 s compiled -> 4.8 s with both.
NaNs and Box-Cox use scikit-learn. TABCTX_FAST_PREPROCESSING=0 turns it off.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 CHANGELOG.md                       |   1 +
 README.md                          |   4 +
 ROADMAP.md                         |  10 +-
 benchmarks/bench_overlap.py        |   3 +
 benchmarks/eval_kv_precision.py    |  15 ++-
 benchmarks/profile_fit_predict.py  |  13 +-
 src/tabctx/backends/tabicl.py      |  16 ++-
 src/tabctx/backends/yeo_johnson.py | 207 +++++++++++++++++++++++++++++
 src/tabctx/serve/factory.py        |   8 ++
 tests/unit/test_serve_factory.py   |   8 ++
 tests/unit/test_yeo_johnson.py     |  74 +++++++++++
 11 files changed, 347 insertions(+), 12 deletions(-)
 create mode 100644 src/tabctx/backends/yeo_johnson.py
 create mode 100644 tests/unit/test_yeo_johnson.py

diff --git a/CHANGELOG.md b/CHANGELOG.md
index c6930b5..3b5c2ae 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -15,6 +15,7 @@
 - ROADMAP: multi-GPU and CUDA-level track (#30-#37), PyTorch-level first.
 
 ### Changed
+- **TabICL's power transforms are fitted on the GPU** (`TABCTX_FAST_PREPROCESSING`, default on; `backends/yeo_johnson.py`). tabicl normalizes each ensemble member with scikit-learn's Yeo-Johnson PowerTransformer, which fits one lambda per column with scipy -- 80%+ of a fit's CPU half. A vectorized fit (same objective, bounds and tolerance; grid then golden-section search, all columns at once) runs in 0.08-0.25 s instead of 0.8-4.4 s on an A100 host; lambdas agree with scikit-learn's within one float32 ulp, accuracy/AUC/R^2 are unchanged on the evaluation datasets, and NaN inputs or Box-Cox still use scikit-learn. With compiled blocks, an isolated 40k x 100 fit takes 4.8 s (was 11.2 s).
 - **TabICL's attention blocks are compiled** (#36, `TABCTX_COMPILE`, default on; `backends/tabicl_compile.py`). Inductor fuses the memory-bound layer norms, residual adds and GELU MLP around each attention kernel: on an A100 the fit's GPU half is 1.7-2.2x faster (10k x 50: 0.91 -> 0.52 s; 40k x 100: 7.65 -> 4.30 s; 20k x 200: 6.51 -> 2.92 s) and its peak ~20% lower (40k x 100: 33.1 -> 27.0 GiB); 100-row predicts 20-35% faster. Accuracy, AUC and R^2 are unchanged on the evaluation datasets, and predicts beside a fit stay correct. Ray Serve, 6 mixed users on 4x A100: 5.78 -> 6.99 predicts/s, small-fit p95 1.48 -> 0.95 s. tabicl's ssmax log(n) is now computed symbolically inside compiled code, so a new table size doesn't recompile. Only tabicl's fp16 path (1,024+ rows) is compiled; its common graphs (~50) compile in `warmup()` at replica start (~130 s per worker, workers in parallel); shapes it doesn't cover (very large tables) compile on first use, ~10 s once per worker. Compile errors fall back to eager, with a warning after warmup if nothing compiled.
 - **Host memory is admitted by estimated bytes, node-wide** (`memory/host.py`). Each backend estimates a call's host use (TabICL ~5 KB per cell when it offloads stage outputs, else nothing; TabPFN ~0.6 KB per cell) and the guard admits calls against one budget shared by every worker process, first come first served (a large fit can't be starved by modest ones), with at most `slots` offloading calls at once (each budgets 1/slots of free host memory for staging). Calls under 1 GiB never wait. The old guard gave every call over 3M cells a whole 48 GiB slot, which serialized fits that use a few GiB (40k x 100 TabICL fits took p50 195 s under concurrent load on the 4x L4 rehearsal). A worker that dies holding host memory has it returned by its parent, and a stuck accounting lock fails calls with a retryable 503 instead of hanging. **`/readyz` `host_memory` keys changed**: `large_fits_active`/`large_fits_waiting` are now `calls_active`/`calls_waiting`, plus `reserved_bytes` and `calls_offloading`.
 - **TabICL fits prepare their CPU half before taking the device**: preprocessing (~30-37% of a fit) overlaps another fit's GPU half -- back-to-back fits 17-26% faster on an A100. At most one fit per device prepares ahead, so queued fits don't each hold preprocessed copies in host memory.
diff --git a/README.md b/README.md
index 54a8dd3..ebd02a6 100644
--- a/README.md
+++ b/README.md
@@ -402,6 +402,10 @@ pip install "tabctx[dev]"        # + test dependencies
   memory, and accuracy is unchanged. The compile happens mostly once per
   worker at startup (~130 s); a call shape the startup didn't cover (very
   large tables, for one) compiles on first use, ~10 s once per worker.
+- **GPU preprocessing** (`TABCTX_FAST_PREPROCESSING=1`, the default):
+  TabICL's per-column Yeo-Johnson fits, most of a fit's CPU time, run
+  vectorized on the GPU (14x faster on a 40k x 100 table) with lambdas
+  matching scikit-learn's to float32 precision.
 - **Same-context coalescing** (`CoalescingPredictor`): concurrent
   requests against one cached context are packed into a single GPU call
   within a ~5ms window and split back per caller, amortizing per-call
diff --git a/ROADMAP.md b/ROADMAP.md
index 1eec921..77598ef 100644
--- a/ROADMAP.md
+++ b/ROADMAP.md
@@ -71,10 +71,12 @@ no FP8).
    half is prepared before taking the device (overlaps another fit's GPU
    half: back-to-back fits 17-26% faster), and tabicl's per-call
    log(n) host-to-device syncs are cached (bit-identical; 10k x 50 fits
-   16% faster), and the attention blocks are compiled (#36, below). Next
-   from it: a GPU Yeo-Johnson fit for isolated-fit latency (~3 ms vs
-   2.9 s at 40k x 100; not bit-identical to sklearn's lambdas, so behind
-   an accuracy study). TabPFN-3.5's fit is the same shape: RMS norm 33% of
+   16% faster), and the attention blocks are compiled (#36, below). The
+   CPU half's Yeo-Johnson fits are vectorized on the GPU
+   (`backends/yeo_johnson.py`: 40k x 100 3.38 -> 0.25 s, lambdas within
+   one float32 ulp of scikit-learn's, accuracy unchanged on the
+   evaluation sets). An isolated 40k x 100 fit: 11.2 -> 7.7 s (compiled)
+   -> 4.8 s (both). TabPFN-3.5's fit is the same shape: RMS norm 33% of
    its CUDA time, flash attention 26%, copies 16%.
 7. **Per-device calibration grids** ([#37](https://github.com/VectorInstitute/tabctx/issues/37), TabPFN done): TabPFN-3.5 has its
    own requirement grid (A100, 20 shapes): it, too, sizes working memory
diff --git a/benchmarks/bench_overlap.py b/benchmarks/bench_overlap.py
index 885e59e..89080d2 100644
--- a/benchmarks/bench_overlap.py
+++ b/benchmarks/bench_overlap.py
@@ -56,6 +56,9 @@ def main() -> int:
     p.add_argument("--test-rows", type=int, default=100)
     p.add_argument("--out", default=None)
     p.add_argument("--compile", action="store_true", help="compiled TabICL blocks")
+    p.add_argument(
+        "--fast-preprocessing", action="store_true", help="vectorized Yeo-Johnson"
+    )
     args = p.parse_args()
     device = torch.device(args.device)
 
diff --git a/benchmarks/eval_kv_precision.py b/benchmarks/eval_kv_precision.py
index cb9f2b6..0014382 100644
--- a/benchmarks/eval_kv_precision.py
+++ b/benchmarks/eval_kv_precision.py
@@ -12,9 +12,10 @@
 
 Usage: eval_kv_precision.py --device cuda:0 --out kv-precision-a100.json
 
-With --compile, the same with compiled attention blocks (#32/#36): run it
-once with and once without and compare the metrics (compilation patches
-tabicl process-wide, so eager and compiled can't share a process).
+With --compile (compiled attention blocks, #32/#36) or
+--fast-preprocessing (vectorized Yeo-Johnson, #31): run once with and once
+without and compare the metrics (both patch tabicl process-wide, so they
+can't share a process with the baseline).
 """
 
 from __future__ import annotations
@@ -85,11 +86,17 @@ def main() -> int:
     p.add_argument("--out", default=None)
     p.add_argument("--precisions", nargs="+", default=list(PRECISIONS))
     p.add_argument("--compile", action="store_true", help="compiled TabICL blocks")
+    p.add_argument(
+        "--fast-preprocessing", action="store_true", help="vectorized Yeo-Johnson"
+    )
     args = p.parse_args()
 
     backends = {
         prec: TabICLBackend(
-            device=args.device, kv_precision=prec, compile_blocks=args.compile
+            device=args.device,
+            kv_precision=prec,
+            compile_blocks=args.compile,
+            fast_preprocessing=args.fast_preprocessing,
         )
         for prec in args.precisions
     }
diff --git a/benchmarks/profile_fit_predict.py b/benchmarks/profile_fit_predict.py
index e7f4d2e..aebf778 100644
--- a/benchmarks/profile_fit_predict.py
+++ b/benchmarks/profile_fit_predict.py
@@ -51,13 +51,22 @@ def make_table(rows: int, features: int, task: str, seed: int):
 
 
 def build_backend(
-    kind: str, device: str, cache_mode: str, compile_blocks: bool = False
+    kind: str,
+    device: str,
+    cache_mode: str,
+    compile_blocks: bool = False,
+    fast_preprocessing: bool = False,
 ):
     if kind == "tabicl":
         from tabctx.backends.tabicl import TabICLBackend
 
         kv: bool | str = False if cache_mode == "off" else cache_mode
-        return TabICLBackend(device=device, kv_cache=kv, compile_blocks=compile_blocks)
+        return TabICLBackend(
+            device=device,
+            kv_cache=kv,
+            compile_blocks=compile_blocks,
+            fast_preprocessing=fast_preprocessing,
+        )
     from tabctx.backends.tabpfn import TabPFNBackend
 
     return TabPFNBackend(device=device, cache_mode=cache_mode)
diff --git a/src/tabctx/backends/tabicl.py b/src/tabctx/backends/tabicl.py
index e2ec377..23acd7c 100644
--- a/src/tabctx/backends/tabicl.py
+++ b/src/tabctx/backends/tabicl.py
@@ -56,6 +56,7 @@
 
 import numpy as np
 
+from tabctx.backends import yeo_johnson
 from tabctx.backends._cuda import (
     QuiescedMeasurement,
     device_scope,
@@ -213,6 +214,7 @@ def __init__(
         offload_dir: str | None = None,
         cpu_safety_factor: float = 0.6,
         compile_blocks: bool = False,
+        fast_preprocessing: bool = False,
     ) -> None:
         """kv_cache: passed through to TabICL ("kv", "repr", True, or
         False). tabicl's own default is False, which makes every
@@ -228,7 +230,12 @@ def __init__(
 
         compile_blocks: torch.compile tabicl's attention blocks (CUDA only;
         see backends/tabicl_compile.py): 1.7-2.2x faster GPU half, ~20%
-        lower peak, paid for by a one-time compile in warmup()."""
+        lower peak, paid for by a one-time compile in warmup().
+
+        fast_preprocessing: fit the ensemble's Yeo-Johnson power transforms
+        for all columns at once, on the GPU (backends/yeo_johnson.py)
+        instead of scikit-learn's per-column scipy search -- 80%+ of a
+        fit's CPU half. Lambdas agree to ~1e-6; not bit-identical."""
         import torch
 
         self._device = device or ("cuda" if torch.cuda.is_available() else "cpu")
@@ -268,6 +275,8 @@ def __init__(
         _install_restore_hook(TabICLRegressor)
         if not install_cached_logn():
             logger.warning("tabicl.ssmax._logn not found: per-call GPU syncs remain")
+        if fast_preprocessing and not yeo_johnson.install():
+            logger.warning("tabicl preprocessing not patchable: scipy power transforms")
         self._compiled = False
         if compile_blocks and is_cuda(self._device):
             from tabctx.backends.tabicl_compile import install_compiled_blocks
@@ -332,8 +341,11 @@ def prepare_fit(self, X: ArrayLike, y: ArrayLike, task: Task) -> _PreparedFit:
         model = self._make_estimator(task)
         kv_cache = model.kv_cache
         model.kv_cache = False
+        # Power transforms fit on this backend's GPU when it is one (see
+        # backends/yeo_johnson.py; a no-op unless fast_preprocessing).
+        transforms_on = self._device if is_cuda(self._device) else None
         try:
-            with device_scope(self._device):
+            with device_scope(self._device), yeo_johnson.on_device(transforms_on):
                 model.fit(X, y)
         finally:
             model.kv_cache = kv_cache
diff --git a/src/tabctx/backends/yeo_johnson.py b/src/tabctx/backends/yeo_johnson.py
new file mode 100644
index 0000000..b403354
--- /dev/null
+++ b/src/tabctx/backends/yeo_johnson.py
@@ -0,0 +1,207 @@
+"""A vectorized Yeo-Johnson PowerTransformer for TabICL's preprocessing (#31).
+
+TabICL normalizes every ensemble member's features with scikit-learn's
+PowerTransformer(method="yeo-johnson"), which fits one lambda per column
+with scipy's bounded Brent search -- a Python loop over columns, dozens of
+passes over each. On an A100 host that was 80%+ of a fit's CPU half
+(40,000 x 100: 3.2 of 3.9 s; 20,000 x 200: 4.0 of 4.6 s), and with the
+GPU half compiled (tabicl_compile.py) the CPU half became most of an
+isolated fit's latency.
+
+``FastPowerTransformer`` fits all columns at once with torch, on the
+backend's GPU when one is set (``on_device``): the same objective
+(scipy's yeojohnson_llf), the same per-column bounds (scipy's
+yeojohnson_normmax, which keep the transform from overflowing the data's
+dtype), and the same tolerance -- a coarse grid over the bounds, then
+golden-section refinement of the best bracket. The transform and the
+standardization are scikit-learn's (StandardScaler), so the fitted object
+is an ordinary PowerTransformer to everything downstream (transform,
+pickling for spill).
+
+It is not bit-identical: lambdas agree with scipy's to ~1e-6 (they are
+stored in the data's dtype, float32, anyway), and a column whose
+likelihood has two local optima may get the better one where Brent stops
+at the first. Anything outside the fast path -- Box-Cox, non-finite
+values, NaNs -- runs scikit-learn's implementation unchanged.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import math
+import threading
+from collections.abc import Iterator
+from typing import Any
+
+import numpy as np
+from sklearn.preprocessing import PowerTransformer, StandardScaler
+
+# scipy's optimize.brent / fminbound tolerance (yeojohnson_normmax).
+_XTOL = 1.48e-08
+_GRID_POINTS = 64
+_GOLDEN = (math.sqrt(5.0) - 1.0) / 2.0
+
+_local = threading.local()
+
+
+@contextlib.contextmanager
+def on_device(device: str | None) -> Iterator[None]:
+    """Fit lambdas on `device` (e.g. "cuda:1") within this block, on this
+    thread; CPU otherwise."""
+    previous = getattr(_local, "device", None)
+    _local.device = device
+    try:
+        yield
+    finally:
+        _local.device = previous
+
+
+def _bounds(x_abs_max: np.ndarray, any_neg: np.ndarray, all_neg: np.ndarray, dtype):
+    """scipy's yeojohnson_normmax bounds, per column: lambdas for which
+    values up to 20x the largest observed transform without over- or
+    underflowing half of `dtype`'s exponent range."""
+    finfo = np.finfo(dtype)
+    log1p_max_x = np.log1p(20 * x_abs_max)
+    log_eps = np.log(finfo.eps)
+    log_tiny = (np.log(finfo.tiny) - log_eps) / 2
+    log_max = (np.log(finfo.max) + log_eps) / 2
+    with np.errstate(divide="ignore"):
+        lb = log_tiny / log1p_max_x
+        ub = log_max / log1p_max_x
+    neg_lb, neg_ub = 2 - ub, 2 - lb
+    lb_out = np.where(all_neg, neg_lb, np.where(any_neg, np.maximum(neg_lb, lb), lb))
+    ub_out = np.where(all_neg, neg_ub, np.where(any_neg, np.minimum(neg_ub, ub), ub))
+    return lb_out, ub_out
+
+
+def _transform(x: Any, lam: Any) -> Any:
+    """Yeo-Johnson of each column of x by its lambda (torch, broadcast)."""
+    import torch
+
+    pos = x >= 0
+    lp = torch.log1p(x.abs())
+    eps = torch.finfo(x.dtype).eps
+    near0 = lam.abs() < eps
+    near2 = (lam - 2).abs() < eps
+    safe = torch.where(near0, torch.ones_like(lam), lam)
+    safe2 = torch.where(near2, torch.ones_like(lam), 2 - lam)
+    pos_out = torch.where(near0, lp, torch.expm1(lam * lp) / safe)
+    neg_out = torch.where(near2, -lp, -torch.expm1((2 - lam) * lp) / safe2)
+    return torch.where(pos, pos_out, neg_out)
+
+
+def _neg_llf(x: Any, lam: Any, sum_signed_log: Any) -> Any:
+    """scipy's -yeojohnson_llf for every column at once (inf variance and
+    zero variance count as rejected, as scipy's _neg_llf does)."""
+    import torch
+
+    n = x.shape[0]
+    var = _transform(x, lam).var(dim=0, unbiased=False)
+    llf = -n / 2 * torch.log(var) + (lam - 1) * sum_signed_log
+    return torch.where(torch.isfinite(llf), -llf, torch.full_like(llf, math.inf))
+
+
+def fit_lambdas(X: np.ndarray, device: str | None = None) -> np.ndarray:
+    """The maximum-likelihood Yeo-Johnson lambda of each column of finite X."""
+    import torch
+
+    x = torch.as_tensor(np.asarray(X, dtype=np.float64), device=device or "cpu")
+    x_abs_max = x.abs().max(dim=0).values.cpu().numpy()
+    any_neg = (x < 0).any(dim=0).cpu().numpy()
+    all_neg = (x < 0).all(dim=0).cpu().numpy()
+    lb_np, ub_np = _bounds(x_abs_max, any_neg, all_neg, X.dtype)
+    lb = torch.as_tensor(lb_np, device=x.device, dtype=torch.float64)
+    ub = torch.as_tensor(ub_np, device=x.device, dtype=torch.float64)
+    sum_signed_log = (torch.sign(x) * torch.log1p(x.abs())).sum(dim=0)
+
+    # Coarse grid over each column's bounds: the best point and its
+    # neighbours bracket the optimum.
+    steps = torch.linspace(0, 1, _GRID_POINTS, device=x.device, dtype=torch.float64)
+    width = ub - lb
+    values = torch.stack(
+        [_neg_llf(x, lb + t * width, sum_signed_log) for t in steps]
+    )  # [grid, columns]
+    best = values.argmin(dim=0)
+    step = width / (_GRID_POINTS - 1)
+    a = torch.maximum(lb, lb + (best - 1).clamp(min=0) * step)
+    b = torch.minimum(ub, lb + (best + 1).clamp(max=_GRID_POINTS - 1) * step)
+
+    # Golden-section search within each bracket, to scipy's tolerance.
+    c = b - _GOLDEN * (b - a)
+    d = a + _GOLDEN * (b - a)
+    fc = _neg_llf(x, c, sum_signed_log)
+    fd = _neg_llf(x, d, sum_signed_log)
+    iterations = math.ceil(
+        math.log(_XTOL / max(float((b - a).max()), _XTOL)) / math.log(_GOLDEN)
+    )
+    for _ in range(max(iterations, 0)):
+        left = fc < fd  # the minimum is in [a, d]
+        b = torch.where(left, d, b)
+        a = torch.where(left, a, c)
+        d_new = torch.where(left, c, a + _GOLDEN * (b - a))
+        c_new = torch.where(left, b - _GOLDEN * (b - a), d)
+        f_new = _neg_llf(x, torch.where(left, c_new, d_new), sum_signed_log)
+        fc, fd = torch.where(left, f_new, fd), torch.where(left, fc, f_new)
+        c, d = c_new, d_new
+    return ((a + b) / 2).cpu().numpy()
+
+
+class FastPowerTransformer(PowerTransformer):
+    """PowerTransformer whose Yeo-Johnson fit is vectorized (module
+    docstring); everything else is scikit-learn's."""
+
+    def _fit(self, X, y=None, force_transform=False):
+        if self.method != "yeo-johnson":
+            return super()._fit(X, y, force_transform)
+        X = self._check_input(X, in_fit=True, check_positive=True)
+        if not np.all(np.isfinite(X)):  # NaNs: scikit-learn drops them per column
+            return super()._fit(X, y, force_transform)
+        if not self.copy and not force_transform:
+            X = X.copy()
+        n_samples = X.shape[0]
+        mean = np.mean(X, axis=0, dtype=np.float64)
+        var = np.var(X, axis=0, dtype=np.float64)
+        constant = _constant_features(var, mean, n_samples) | np.all(X == 0, axis=0)
+        lambdas = np.ones(X.shape[1], dtype=X.dtype)
+        varying = np.flatnonzero(~constant)
+        if varying.size:
+            device = getattr(_local, "device", None)
+            lambdas[varying] = fit_lambdas(X[:, varying], device)
+        self.lambdas_ = lambdas
+        if self.standardize or force_transform:
+            import torch
+
+            device = getattr(_local, "device", None) or "cpu"
+            for start in range(0, varying.size, 256):  # bound device memory
+                cols = varying[start : start + 256]
+                x = torch.as_tensor(X[:, cols].astype(np.float64), device=device)
+                lam = torch.as_tensor(lambdas[cols].astype(np.float64), device=device)
+                X[:, cols] = _transform(x, lam).cpu().numpy().astype(X.dtype)
+        if self.standardize:
+            self._scaler = StandardScaler(copy=False).set_output(transform="default")
+            if force_transform:
+                X = self._scaler.fit_transform(X)
+            else:
+                self._scaler.fit(X)
+        return X
+
+
+def _constant_features(var: np.ndarray, mean: np.ndarray, n_samples: int) -> np.ndarray:
+    """scikit-learn's test for a constant column (variance within rounding
+    of zero), which PowerTransformer leaves untransformed (lambda 1)."""
+    eps = np.finfo(np.float64).eps
+    upper_bound = n_samples * eps * var + (n_samples * mean * eps) ** 2
+    return var <= upper_bound
+
+
+def install() -> bool:
+    """Use FastPowerTransformer in tabicl's preprocessing. Idempotent;
+    False if tabicl's preprocessing no longer imports PowerTransformer."""
+    try:
+        from tabicl._sklearn import preprocessing
+    except ImportError:
+        return False
+    if getattr(preprocessing, "PowerTransformer", None) is None:
+        return False
+    preprocessing.PowerTransformer = FastPowerTransformer
+    return True
diff --git a/src/tabctx/serve/factory.py b/src/tabctx/serve/factory.py
index b9b2690..4b15062 100644
--- a/src/tabctx/serve/factory.py
+++ b/src/tabctx/serve/factory.py
@@ -114,6 +114,7 @@
 DEVICES_ENV_VAR = "TABCTX_DEVICES"
 WORKER_PROCESSES_ENV_VAR = "TABCTX_WORKER_PROCESSES"
 COMPILE_ENV_VAR = "TABCTX_COMPILE"
+FAST_PREPROCESSING_ENV_VAR = "TABCTX_FAST_PREPROCESSING"
 
 logger = logging.getLogger(__name__)
 
@@ -144,6 +145,9 @@ class ServeSettings:
     # torch.compile the models' hot blocks where measured to help (TabICL's
     # attention blocks: backends/tabicl_compile.py).
     compile_models: bool = True
+    # Vectorized (GPU) Yeo-Johnson fits in TabICL's preprocessing
+    # (backends/yeo_johnson.py).
+    fast_preprocessing: bool = True
 
     @classmethod
     def from_env(cls) -> ServeSettings:
@@ -234,6 +238,9 @@ def from_env(cls) -> ServeSettings:
             offload_dir=offload_dir,
             worker_processes=_env_flag(WORKER_PROCESSES_ENV_VAR, cls.worker_processes),
             compile_models=_env_flag(COMPILE_ENV_VAR, cls.compile_models),
+            fast_preprocessing=_env_flag(
+                FAST_PREPROCESSING_ENV_VAR, cls.fast_preprocessing
+            ),
         )
 
 
@@ -511,6 +518,7 @@ def _build_real_backend(
             # headroom (see memory/host.py).
             cpu_safety_factor=0.6 / host_slots,
             compile_blocks=settings.compile_models,
+            fast_preprocessing=settings.fast_preprocessing,
         )
     # Weights load at replica start, not inside (and charged to) the first
     # tenant's fit.
diff --git a/tests/unit/test_serve_factory.py b/tests/unit/test_serve_factory.py
index ab7274c..b39e047 100644
--- a/tests/unit/test_serve_factory.py
+++ b/tests/unit/test_serve_factory.py
@@ -482,3 +482,11 @@ def test_compile_flag_from_env(monkeypatch):
     monkeypatch.setenv(COMPILE_ENV_VAR, "maybe")
     with pytest.raises(ValueError, match="TABCTX_COMPILE"):
         ServeSettings.from_env()
+
+
+def test_fast_preprocessing_flag_from_env(monkeypatch):
+    from tabctx.serve.factory import FAST_PREPROCESSING_ENV_VAR
+
+    assert ServeSettings.from_env().fast_preprocessing
+    monkeypatch.setenv(FAST_PREPROCESSING_ENV_VAR, "0")
+    assert not ServeSettings.from_env().fast_preprocessing
diff --git a/tests/unit/test_yeo_johnson.py b/tests/unit/test_yeo_johnson.py
new file mode 100644
index 0000000..280d304
--- /dev/null
+++ b/tests/unit/test_yeo_johnson.py
@@ -0,0 +1,74 @@
+"""The vectorized Yeo-Johnson PowerTransformer (backends/yeo_johnson.py)
+against scikit-learn's, on CPU torch."""
+
+import numpy as np
+import pytest
+
+pytest.importorskip("torch")
+from sklearn.preprocessing import PowerTransformer  # noqa: E402
+
+from tabctx.backends.yeo_johnson import FastPowerTransformer, fit_lambdas  # noqa: E402
+
+
+def columns(n=4000, seed=0):
+    rng = np.random.default_rng(seed)
+    cols = [
+        rng.normal(size=n),
+        rng.exponential(size=n),
+        -3 * rng.exponential(size=n),  # all negative
+        rng.lognormal(sigma=1.5, size=n),
+        rng.standard_t(3, size=n),
+        rng.uniform(-1, 1, n) ** 3,
+        np.zeros(n),  # constant: lambda 1, untouched
+    ]
+    X = np.stack(cols, 1)
+    X[:, :-1] = (X[:, :-1] - X[:, :-1].mean(0)) / X[:, :-1].std(0)
+    return X.astype(np.float32)
+
+
+@pytest.mark.parametrize("seed", [0, 1, 2])
+def test_matches_scikit_learn(seed):
+    X = columns(seed=seed)
+    ref = PowerTransformer(method="yeo-johnson").fit(X)
+    fast = FastPowerTransformer(method="yeo-johnson").fit(X)
+    assert fast.lambdas_.dtype == ref.lambdas_.dtype
+    np.testing.assert_allclose(fast.lambdas_, ref.lambdas_, atol=1e-5)
+    np.testing.assert_allclose(fast.transform(X), ref.transform(X), atol=1e-4)
+    assert fast.lambdas_[-1] == 1.0
+
+
+def test_fit_transform_matches_too():
+    X = columns()
+    ref = PowerTransformer(method="yeo-johnson").fit_transform(X.copy())
+    fast = FastPowerTransformer(method="yeo-johnson").fit_transform(X.copy())
+    np.testing.assert_allclose(fast, ref, atol=1e-4)
+
+
+def test_nans_and_box_cox_use_scikit_learn():
+    X = columns()
+    X[3, 1] = np.nan
+    ref = PowerTransformer(method="yeo-johnson").fit(X)
+    fast = FastPowerTransformer(method="yeo-johnson").fit(X)
+    np.testing.assert_array_equal(fast.lambdas_, ref.lambdas_)  # same code path
+    positive = np.abs(columns()[:, :3]) + 0.1
+    np.testing.assert_array_equal(
+        FastPowerTransformer(method="box-cox").fit(positive).lambdas_,
+        PowerTransformer(method="box-cox").fit(positive).lambdas_,
+    )
+
+
+def test_lambdas_stay_within_scipys_overflow_bounds():
+    rng = np.random.default_rng(3)
+    X = np.stack([rng.lognormal(sigma=4, size=2000), rng.normal(size=2000)], 1)
+    lam = fit_lambdas(X.astype(np.float32))
+    out = FastPowerTransformer(method="yeo-johnson").fit_transform(X.astype(np.float32))
+    assert np.all(np.isfinite(lam)) and np.all(np.isfinite(out))
+
+
+def test_the_fitted_transformer_pickles_as_a_power_transformer():
+    import pickle
+
+    fast = FastPowerTransformer(method="yeo-johnson").fit(columns())
+    again = pickle.loads(pickle.dumps(fast))
+    assert isinstance(again, PowerTransformer)
+    np.testing.assert_array_equal(again.lambdas_, fast.lambdas_)

From 064d2b856a71faec11d4e8c8d2b94c91eebab94a Mon Sep 17 00:00:00 2001
From: Amrit Krishnan <amrit110@gmail.com>
Date: Sun, 18 Oct 2026 01:16:53 -0400
Subject: [PATCH 33/39] fix: review of placement/events and the GPU Yeo-Johnson

Yeo-Johnson (backends/yeo_johnson.py):
- The device section runs under the backend's measurement lock and frees
  everything before leaving it: preparation runs beside another fit's GPU
  half, and its temporaries in that fit's before/after allocator readings
  understated the context (the unsafe direction).
- Columns are fitted in chunks bounded to 512 MiB of temporaries (100k x
  300: 519 MiB peak), the allocator cache is returned after, and a CUDA
  out-of-memory falls back to scikit-learn instead of failing the fit.
- The fast path is per backend (on_device sets it per thread): install()
  no longer changes other backends in the process.
- Fitted transformers pickle as scikit-learn's PowerTransformer, so spill
  blobs name nothing from tabctx.
- Values beyond any overflow-safe lambda get lambda 1 explicitly.

Residency events: fit_predict's throwaway contexts aren't reported; an
evict failing on one device is logged and skipped (no false "removed", the
others still cleaned); MultiDeviceEngine.residency() and
/v1/tabctx/residency give each device's resident and spilled ids to
reconcile against, since events are advisory.

Tests: randomized equivalence of the default policy with the previous
rule; events for fit_predict, failing policies and evicts; YJ under
on_device, guard, chunking, OOM fallback, pickling, float64, degenerate
bounds, tabicl install.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 CHANGELOG.md                        |   4 +-
 src/tabctx/backends/tabicl.py       |  28 ++++--
 src/tabctx/backends/yeo_johnson.py  | 146 ++++++++++++++++++++++------
 src/tabctx/events.py                |   7 +-
 src/tabctx/multidevice.py           |  68 ++++++++++---
 src/tabctx/serve/app.py             |  16 ++-
 src/tabctx/serve/workers.py         |   6 ++
 tests/unit/test_placement_events.py |  81 +++++++++++++++
 tests/unit/test_yeo_johnson.py      | 117 ++++++++++++++++++++--
 9 files changed, 407 insertions(+), 66 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3b5c2ae..2688db6 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,14 +8,14 @@
 - **One worker process per GPU** (#39, `TABCTX_WORKER_PROCESSES`, default on for multi-GPU replicas). Fits are host-bound, so threads of one process driving several GPUs contend for the GIL: on 4x A100, four concurrent fits ran at 2.1-3.1x aggregate from threads and 3.96-4.3x from processes (`benchmarks/bench_scaling.py`). `serve/workers.py` runs each device's engine in its own process behind a `ProcessShard` proxy; the router reads versioned residency/headroom snapshots that every reply carries, so its decisions stay lock-free. A dead worker fails its calls, is skipped by placement, and fails `/readyz` and Ray Serve's `check_health` so the replica restarts. Nothing depends on Ray. Under Ray Serve on 4x A100 with 6 mixed users: 2.58 -> 5.0 predicts/s, small-fit p50 1.6 -> 0.6 s.
 - **Light predicts run beside a fit** on the same GPU (#39). A predict against a resident context with at most 2,048 rows no longer waits for a fit on its device (it waited up to 266 s on the L4 rehearsal): ~0.1 s p50 during a 40k x 100 fit on an A100, at ~15% fit slowdown. Restores and large predicts still wait; contexts in use are pinned against evict-ahead (and evicted with predicts held off when nothing else makes room); a fit that runs out of memory beside a predict retries alone; allocator measurements are taken with predicts quiesced, and memory freed during a fit is released only after it has measured itself.
 - **TabPFN-3.5 requirement calibration** (A100, 20 shapes up to 200k x 10 and 30k x 1,000). TabPFN also sizes its working memory from what is free (20k x 100 peaks at 7.3 GiB unconstrained but runs within 2x its time on 1.1 GiB), so admission bounds resident + working floor for it too, instead of TabICL's peak formula: up to 100k rows at 100 features and 80k at 1,000 on an A40 (was 5k and 500 on an L4). Its per-row resident rate grows with features, which `ResidentRateModel(width_dependent=True)` respects.
-- **Pluggable placement and residency events** (#39). The multi-device router asks a `PlacementPolicy` (`placement.py`) which device takes a fit, from a lock-free view of each device (headroom, the fit's estimated need there, activity, liveness); the previous rule is the default `LeastActiveThatFits`, which now also skips dead workers. The router emits `ResidencyEvent`s (`events.py`) when it places a context and when it removes a copy, to any sink; `/v1/tabctx/residency` serves the latest. A platform that co-schedules LLMs and tabular models can supply its own policy and follow placements without polling.
+- **Pluggable placement and residency events** (#39). The multi-device router asks a `PlacementPolicy` (`placement.py`) which device takes a fit, from a lock-free view of each device (headroom, the fit's estimated need there, activity, liveness); the previous rule is the default `LeastActiveThatFits`, which now also skips dead workers. The router emits `ResidencyEvent`s (`events.py`) when it places a context and when it removes a copy, to any sink; `/v1/tabctx/residency` serves the latest. A platform that co-schedules LLMs and tabular models can supply its own policy and follow placements without polling; events are advisory (contexts dropped by capacity eviction without spill aren't reported), so `/v1/tabctx/residency` also returns each device's current resident and spilled ids to reconcile against. `fit_predict`'s throwaway contexts aren't reported, and an evict that fails on one device (a dead worker) no longer stops the others.
 - `benchmarks/bench_scaling.py` (threads vs processes across GPUs), `bench_overlap.py` (predict latency and correctness during a fit), `gke/serve_workload.sh` (Ray Serve + workload inside the devbox).
 - `benchmarks/bench_workload.py`: concurrent users looping fit-then-predict over a small/medium/large table mix (large ones via CSV upload), with per-class latency percentiles, 413s, errors and the per-device spread.
 - `benchmarks/sweep_admission.py`: the table-size sweep (rows up to the first 413 per feature count, fit and predict timed).
 - ROADMAP: multi-GPU and CUDA-level track (#30-#37), PyTorch-level first.
 
 ### Changed
-- **TabICL's power transforms are fitted on the GPU** (`TABCTX_FAST_PREPROCESSING`, default on; `backends/yeo_johnson.py`). tabicl normalizes each ensemble member with scikit-learn's Yeo-Johnson PowerTransformer, which fits one lambda per column with scipy -- 80%+ of a fit's CPU half. A vectorized fit (same objective, bounds and tolerance; grid then golden-section search, all columns at once) runs in 0.08-0.25 s instead of 0.8-4.4 s on an A100 host; lambdas agree with scikit-learn's within one float32 ulp, accuracy/AUC/R^2 are unchanged on the evaluation datasets, and NaN inputs or Box-Cox still use scikit-learn. With compiled blocks, an isolated 40k x 100 fit takes 4.8 s (was 11.2 s).
+- **TabICL's power transforms are fitted on the GPU** (`TABCTX_FAST_PREPROCESSING`, default on; `backends/yeo_johnson.py`). tabicl normalizes each ensemble member with scikit-learn's Yeo-Johnson PowerTransformer, which fits one lambda per column with scipy -- 80%+ of a fit's CPU half. A vectorized fit (same objective, bounds and tolerance; grid then golden-section search, all columns at once) runs in 0.08-0.25 s instead of 0.8-4.4 s on an A100 host; lambdas agree with scikit-learn's within one float32 ulp, accuracy/AUC/R^2 are unchanged on the evaluation datasets, and NaN inputs, Box-Cox and a device out-of-memory fall back to scikit-learn. The device section runs under the measurement lock (a concurrent fit's allocator readings never see it) with temporaries bounded to 512 MiB per column chunk, and fitted transformers pickle as plain scikit-learn objects. With compiled blocks, an isolated 40k x 100 fit takes 4.8 s (was 11.2 s).
 - **TabICL's attention blocks are compiled** (#36, `TABCTX_COMPILE`, default on; `backends/tabicl_compile.py`). Inductor fuses the memory-bound layer norms, residual adds and GELU MLP around each attention kernel: on an A100 the fit's GPU half is 1.7-2.2x faster (10k x 50: 0.91 -> 0.52 s; 40k x 100: 7.65 -> 4.30 s; 20k x 200: 6.51 -> 2.92 s) and its peak ~20% lower (40k x 100: 33.1 -> 27.0 GiB); 100-row predicts 20-35% faster. Accuracy, AUC and R^2 are unchanged on the evaluation datasets, and predicts beside a fit stay correct. Ray Serve, 6 mixed users on 4x A100: 5.78 -> 6.99 predicts/s, small-fit p95 1.48 -> 0.95 s. tabicl's ssmax log(n) is now computed symbolically inside compiled code, so a new table size doesn't recompile. Only tabicl's fp16 path (1,024+ rows) is compiled; its common graphs (~50) compile in `warmup()` at replica start (~130 s per worker, workers in parallel); shapes it doesn't cover (very large tables) compile on first use, ~10 s once per worker. Compile errors fall back to eager, with a warning after warmup if nothing compiled.
 - **Host memory is admitted by estimated bytes, node-wide** (`memory/host.py`). Each backend estimates a call's host use (TabICL ~5 KB per cell when it offloads stage outputs, else nothing; TabPFN ~0.6 KB per cell) and the guard admits calls against one budget shared by every worker process, first come first served (a large fit can't be starved by modest ones), with at most `slots` offloading calls at once (each budgets 1/slots of free host memory for staging). Calls under 1 GiB never wait. The old guard gave every call over 3M cells a whole 48 GiB slot, which serialized fits that use a few GiB (40k x 100 TabICL fits took p50 195 s under concurrent load on the 4x L4 rehearsal). A worker that dies holding host memory has it returned by its parent, and a stuck accounting lock fails calls with a retryable 503 instead of hanging. **`/readyz` `host_memory` keys changed**: `large_fits_active`/`large_fits_waiting` are now `calls_active`/`calls_waiting`, plus `reserved_bytes` and `calls_offloading`.
 - **TabICL fits prepare their CPU half before taking the device**: preprocessing (~30-37% of a fit) overlaps another fit's GPU half -- back-to-back fits 17-26% faster on an A100. At most one fit per device prepares ahead, so queued fits don't each hold preprocessed copies in host memory.
diff --git a/src/tabctx/backends/tabicl.py b/src/tabctx/backends/tabicl.py
index 23acd7c..2fb7ece 100644
--- a/src/tabctx/backends/tabicl.py
+++ b/src/tabctx/backends/tabicl.py
@@ -49,6 +49,7 @@
 
 from __future__ import annotations
 
+import contextlib
 import logging
 import threading
 from dataclasses import dataclass
@@ -205,6 +206,7 @@ class TabICLBackend(QuiescedMeasurement):
     # The exact model id callers select via the API's `model` field --
     # named for the checkpoint actually loaded (tabicl-classifier-v2-*).
     name = "tabicl-v2"
+    _fast_preprocessing = False
 
     def __init__(
         self,
@@ -275,7 +277,8 @@ def __init__(
         _install_restore_hook(TabICLRegressor)
         if not install_cached_logn():
             logger.warning("tabicl.ssmax._logn not found: per-call GPU syncs remain")
-        if fast_preprocessing and not yeo_johnson.install():
+        self._fast_preprocessing = fast_preprocessing and yeo_johnson.install()
+        if fast_preprocessing and not self._fast_preprocessing:
             logger.warning("tabicl preprocessing not patchable: scipy power transforms")
         self._compiled = False
         if compile_blocks and is_cuda(self._device):
@@ -332,20 +335,29 @@ def prepare_fit(self, X: ArrayLike, y: ArrayLike, task: Task) -> _PreparedFit:
         """The CPU half of a fit: validation, label encoding and the
         ensemble's preprocessing (per-feature power transforms, shuffles) --
         tabicl's fit() up to building the kv cache, run with the cache
-        disabled. 30-37% of a fit's time (A100, 2k x 20 to 40k x 100); it
-        needs no device memory, so the engine runs it before taking the
-        device, overlapping another fit's GPU half."""
+        disabled. 30-37% of a fit's time (A100, 2k x 20 to 40k x 100) with
+        scikit-learn's power transforms, ~10% with fast_preprocessing. It
+        needs no device memory beyond the power transforms' bounded,
+        measurement-locked temporaries, so the engine runs it before
+        taking the device, overlapping another fit's GPU half."""
         # Fresh estimator instance every call (load-bearing -- see
         # backends/base.py), but the underlying pretrained nn.Module is
         # shared across fits after the first (see __init__).
         model = self._make_estimator(task)
         kv_cache = model.kv_cache
         model.kv_cache = False
-        # Power transforms fit on this backend's GPU when it is one (see
-        # backends/yeo_johnson.py; a no-op unless fast_preprocessing).
-        transforms_on = self._device if is_cuda(self._device) else None
+        # With fast_preprocessing, power transforms fit on this backend's GPU
+        # (backends/yeo_johnson.py), under the measurement lock: this runs
+        # beside another fit's GPU half, whose readings must not see it.
+        fast = (
+            yeo_johnson.on_device(
+                self._device if is_cuda(self._device) else None, self.quiesced
+            )
+            if self._fast_preprocessing
+            else contextlib.nullcontext()
+        )
         try:
-            with device_scope(self._device), yeo_johnson.on_device(transforms_on):
+            with device_scope(self._device), fast:
                 model.fit(X, y)
         finally:
             model.kv_cache = kv_cache
diff --git a/src/tabctx/backends/yeo_johnson.py b/src/tabctx/backends/yeo_johnson.py
index b403354..5f29052 100644
--- a/src/tabctx/backends/yeo_johnson.py
+++ b/src/tabctx/backends/yeo_johnson.py
@@ -18,6 +18,14 @@
 is an ordinary PowerTransformer to everything downstream (transform,
 pickling for spill).
 
+Sharing the device. Preparation runs outside the engine's device locks,
+beside another fit's GPU half (engine.py). So the GPU section runs under
+the backend's measurement lock (the one QuiescedMeasurement readings
+take; ``on_device(device, guard)``) and frees everything before leaving
+it: a concurrent fit's allocator readings never see these temporaries.
+Columns are processed in chunks of at most ``_CHUNK_BYTES`` of
+temporaries, and a CUDA out-of-memory falls back to scikit-learn.
+
 It is not bit-identical: lambdas agree with scipy's to ~1e-6 (they are
 stored in the data's dtype, float32, anyway), and a column whose
 likelihood has two local optima may get the better one where Brent stops
@@ -28,9 +36,10 @@
 from __future__ import annotations
 
 import contextlib
+import logging
 import math
 import threading
-from collections.abc import Iterator
+from collections.abc import Callable, Iterator
 from typing import Any
 
 import numpy as np
@@ -38,22 +47,32 @@
 
 # scipy's optimize.brent / fminbound tolerance (yeojohnson_normmax).
 _XTOL = 1.48e-08
+# Device temporaries per chunk of columns: ~48 bytes per cell live while
+# evaluating the likelihood (float64 input, log1p, both branches, output).
+_CHUNK_BYTES = 512 * 1024**2
+_BYTES_PER_CELL = 48
 _GRID_POINTS = 64
 _GOLDEN = (math.sqrt(5.0) - 1.0) / 2.0
 
 _local = threading.local()
+logger = logging.getLogger(__name__)
 
 
 @contextlib.contextmanager
-def on_device(device: str | None) -> Iterator[None]:
-    """Fit lambdas on `device` (e.g. "cuda:1") within this block, on this
-    thread; CPU otherwise."""
-    previous = getattr(_local, "device", None)
-    _local.device = device
+def on_device(
+    device: str | None,
+    guard: Callable[[], contextlib.AbstractContextManager[Any]] | None = None,
+) -> Iterator[None]:
+    """Enable the fast fit on this thread within this block: on `device`
+    (e.g. "cuda:1"; CPU when None), with each device section run inside
+    `guard()` (the backend's measurement lock). Outside such a block,
+    FastPowerTransformer is scikit-learn's PowerTransformer."""
+    previous = getattr(_local, "config", None)
+    _local.config = (device, guard or contextlib.nullcontext)
     try:
         yield
     finally:
-        _local.device = previous
+        _local.config = previous
 
 
 def _bounds(x_abs_max: np.ndarray, any_neg: np.ndarray, all_neg: np.ndarray, dtype):
@@ -102,7 +121,19 @@ def _neg_llf(x: Any, lam: Any, sum_signed_log: Any) -> Any:
 
 
 def fit_lambdas(X: np.ndarray, device: str | None = None) -> np.ndarray:
-    """The maximum-likelihood Yeo-Johnson lambda of each column of finite X."""
+    """The maximum-likelihood Yeo-Johnson lambda of each column of finite X,
+    in column chunks bounded by _CHUNK_BYTES of device temporaries."""
+    n_rows = max(1, X.shape[0])
+    per_chunk = max(1, _CHUNK_BYTES // (_BYTES_PER_CELL * n_rows))
+    return np.concatenate(
+        [
+            _fit_lambdas(X[:, start : start + per_chunk], device)
+            for start in range(0, X.shape[1], per_chunk)
+        ]
+    )
+
+
+def _fit_lambdas(X: np.ndarray, device: str | None) -> np.ndarray:
     import torch
 
     x = torch.as_tensor(np.asarray(X, dtype=np.float64), device=device or "cpu")
@@ -143,40 +174,42 @@ def fit_lambdas(X: np.ndarray, device: str | None = None) -> np.ndarray:
         f_new = _neg_llf(x, torch.where(left, c_new, d_new), sum_signed_log)
         fc, fd = torch.where(left, f_new, fd), torch.where(left, fc, f_new)
         c, d = c_new, d_new
-    return ((a + b) / 2).cpu().numpy()
+    lam = ((a + b) / 2).cpu().numpy()
+    # Values so large that no lambda keeps them within the dtype's range on
+    # both signs (scipy raises here): leave the column untransformed.
+    return np.where(lb_np >= ub_np, 1.0, lam)
 
 
 class FastPowerTransformer(PowerTransformer):
-    """PowerTransformer whose Yeo-Johnson fit is vectorized (module
-    docstring); everything else is scikit-learn's."""
+    """PowerTransformer whose Yeo-Johnson fit is vectorized inside an
+    ``on_device`` block (module docstring); scikit-learn's otherwise.
+    Pickles as a plain PowerTransformer: the fitted state is identical,
+    so spilled contexts don't depend on this module."""
 
     def _fit(self, X, y=None, force_transform=False):
-        if self.method != "yeo-johnson":
+        config = getattr(_local, "config", None)
+        if self.method != "yeo-johnson" or config is None:
             return super()._fit(X, y, force_transform)
         X = self._check_input(X, in_fit=True, check_positive=True)
         if not np.all(np.isfinite(X)):  # NaNs: scikit-learn drops them per column
             return super()._fit(X, y, force_transform)
+        device, guard = config
+        try:
+            with guard():
+                lambdas, transformed = self._fit_columns(
+                    X, device, self.standardize or force_transform
+                )
+                _release_device_cache(device)
+        except Exception as e:  # noqa: BLE001 -- only a device OOM is expected
+            if not _is_out_of_memory(e):
+                raise
+            logger.warning("Yeo-Johnson fit out of device memory: scikit-learn's")
+            return super()._fit(X, y, force_transform)
         if not self.copy and not force_transform:
             X = X.copy()
-        n_samples = X.shape[0]
-        mean = np.mean(X, axis=0, dtype=np.float64)
-        var = np.var(X, axis=0, dtype=np.float64)
-        constant = _constant_features(var, mean, n_samples) | np.all(X == 0, axis=0)
-        lambdas = np.ones(X.shape[1], dtype=X.dtype)
-        varying = np.flatnonzero(~constant)
-        if varying.size:
-            device = getattr(_local, "device", None)
-            lambdas[varying] = fit_lambdas(X[:, varying], device)
         self.lambdas_ = lambdas
-        if self.standardize or force_transform:
-            import torch
-
-            device = getattr(_local, "device", None) or "cpu"
-            for start in range(0, varying.size, 256):  # bound device memory
-                cols = varying[start : start + 256]
-                x = torch.as_tensor(X[:, cols].astype(np.float64), device=device)
-                lam = torch.as_tensor(lambdas[cols].astype(np.float64), device=device)
-                X[:, cols] = _transform(x, lam).cpu().numpy().astype(X.dtype)
+        if transformed is not None:
+            X[:] = transformed
         if self.standardize:
             self._scaler = StandardScaler(copy=False).set_output(transform="default")
             if force_transform:
@@ -185,6 +218,53 @@ def _fit(self, X, y=None, force_transform=False):
                 self._scaler.fit(X)
         return X
 
+    @staticmethod
+    def _fit_columns(X: np.ndarray, device: str | None, transform: bool):
+        """Lambdas (lambda 1 for constant columns, as scikit-learn) and, if
+        asked, X transformed by them (constant columns untouched)."""
+        import torch
+
+        n_samples = X.shape[0]
+        mean = np.mean(X, axis=0, dtype=np.float64)
+        var = np.var(X, axis=0, dtype=np.float64)
+        constant = _constant_features(var, mean, n_samples) | np.all(X == 0, axis=0)
+        lambdas = np.ones(X.shape[1], dtype=X.dtype)
+        varying = np.flatnonzero(~constant)
+        if varying.size:
+            lambdas[varying] = fit_lambdas(X[:, varying], device)
+        if not transform:
+            return lambdas, None
+        out = X.copy()
+        per_chunk = max(1, _CHUNK_BYTES // (_BYTES_PER_CELL * max(1, n_samples)))
+        for start in range(0, varying.size, per_chunk):
+            cols = varying[start : start + per_chunk]
+            x = torch.as_tensor(X[:, cols].astype(np.float64), device=device or "cpu")
+            lam = torch.as_tensor(lambdas[cols].astype(np.float64), device=x.device)
+            out[:, cols] = _transform(x, lam).cpu().numpy().astype(X.dtype)
+        return lambdas, out
+
+    def __reduce__(self):
+        # Rebuilt as scikit-learn's class (constructed, then its own
+        # __setstate__): the blob names nothing from tabctx.
+        return (PowerTransformer, (), self.__getstate__())
+
+
+def _is_out_of_memory(error: BaseException) -> bool:
+    try:
+        import torch
+    except ImportError:
+        return False
+    return isinstance(error, torch.cuda.OutOfMemoryError)
+
+
+def _release_device_cache(device: str | None) -> None:
+    """Return the temporaries' cached blocks, so free-memory readings (tabicl's
+    offload decision) don't see them as used."""
+    if device is not None and device.startswith("cuda"):
+        import torch
+
+        torch.cuda.empty_cache()
+
 
 def _constant_features(var: np.ndarray, mean: np.ndarray, n_samples: int) -> np.ndarray:
     """scikit-learn's test for a constant column (variance within rounding
@@ -195,8 +275,10 @@ def _constant_features(var: np.ndarray, mean: np.ndarray, n_samples: int) -> np.
 
 
 def install() -> bool:
-    """Use FastPowerTransformer in tabicl's preprocessing. Idempotent;
-    False if tabicl's preprocessing no longer imports PowerTransformer."""
+    """Use FastPowerTransformer in tabicl's preprocessing. Idempotent, and
+    harmless to other backends: it is scikit-learn's PowerTransformer
+    except inside an ``on_device`` block. False if tabicl's preprocessing
+    no longer imports PowerTransformer."""
     try:
         from tabicl._sklearn import preprocessing
     except ImportError:
diff --git a/src/tabctx/events.py b/src/tabctx/events.py
index 610b103..e8cf50d 100644
--- a/src/tabctx/events.py
+++ b/src/tabctx/events.py
@@ -8,8 +8,11 @@
 counter, a Ray actor, a Dynamo event plane publisher.
 
 Events report the router's decisions, not every cache-internal move:
-capacity eviction to the spill tier and restores stay on their device,
-so a context's device doesn't change until the router says so.
+spill and restore keep a context on its device, so its device doesn't
+change until the router says so. They are advisory: capacity eviction
+without a spill tier (or with a full one) drops a context silently, so
+subscribers reconcile against MultiDeviceEngine.residency() (served at
+/v1/tabctx/residency). fit_predict's throwaway contexts aren't reported.
 
 Sinks run on the request's thread after the change, so they must be
 fast and must not raise (one that does is logged and skipped).
diff --git a/src/tabctx/multidevice.py b/src/tabctx/multidevice.py
index 803a454..6d8df6a 100644
--- a/src/tabctx/multidevice.py
+++ b/src/tabctx/multidevice.py
@@ -44,6 +44,7 @@
 from __future__ import annotations
 
 import contextlib
+import logging
 import threading
 import time
 from collections import deque
@@ -61,6 +62,8 @@
 )
 from tabctx.types import ArrayLike, EngineStats, PredictOutcome, Task
 
+logger = logging.getLogger(__name__)
+
 # How long a started call counts toward a device's recent activity.
 ACTIVITY_WINDOW_S = 30.0
 
@@ -88,6 +91,8 @@ def admission_headroom_bytes(self, backend: str | None = None) -> int: ...
 
     def holds(self, dataset_id: str) -> bool: ...
 
+    def residency_nowait(self) -> tuple[frozenset[str], frozenset[str]]: ...
+
     def has(self, dataset_id: str) -> bool: ...
 
     def stats_nowait(self) -> EngineStats: ...
@@ -261,26 +266,55 @@ def fit(
         backend: str | None = None,
         feature_names: list[str] | None = None,
     ) -> str:
+        return self._fit(X, y, task, dataset_id, backend, feature_names, announce=True)
+
+    def _fit(
+        self,
+        X: ArrayLike,
+        y: ArrayLike,
+        task: Task,
+        dataset_id: str | None,
+        backend: str | None,
+        feature_names: list[str] | None,
+        announce: bool,
+    ) -> str:
+        """fit(); `announce` emits residency events (not for fit_predict's
+        throwaway contexts, which would flood subscribers)."""
         n_features = _validate_fit_input(X, y)
         name = self._shards[0].resolve_backend_name(backend)
         request = FitRequest(len(X), n_features, name, dataset_id)
         if dataset_id is None:
             index = self._place(request)
             resolved = self._fit_on(index, X, y, task, None, name, feature_names)
-            self._emit("placed", resolved, index, name)
+            if announce:
+                self._emit("placed", resolved, index, name)
             return resolved
         with self._id_locks.hold(dataset_id):
             index = self._place(request)
             resolved = self._fit_on(index, X, y, task, dataset_id, name, feature_names)
-            self._emit("placed", resolved, index, name)
+            if announce:
+                self._emit("placed", resolved, index, name)
             # Only now that the new context is in place: drop copies on
             # other devices (a failed fit above leaves the old one serving).
-            for i in self._locate(dataset_id):
-                if i != index:
-                    self._shards[i].evict(dataset_id)
-                    self._emit("removed", dataset_id, i)
+            others = [i for i in self._locate(dataset_id) if i != index]
+            self._evict_from(dataset_id, others, announce)
             return resolved
 
+    def _evict_from(self, dataset_id: str, indices: list[int], announce: bool) -> None:
+        """Evict the id from each device, reporting each removal that
+        happened; a device that fails (a dead worker) is logged and skipped,
+        so the others are still cleaned up."""
+        for i in indices:
+            try:
+                self._shards[i].evict(dataset_id)
+            except Exception:  # noqa: BLE001 -- keep cleaning the others
+                logger.exception(
+                    "evicting %s from %s failed", dataset_id, self._devices[i]
+                )
+                continue
+            if announce:
+                self._emit("removed", dataset_id, i)
+
     def _emit(
         self, kind: EventKind, dataset_id: str, index: int, backend: str | None = None
     ) -> None:
@@ -336,17 +370,29 @@ def fit_predict(
         task: Task = "classification",
         return_proba: bool = False,
     ) -> PredictOutcome:
-        dataset_id = self.fit(X, y, task)
+        dataset_id = self._fit(X, y, task, None, None, None, announce=False)
         try:
             return self.predict(dataset_id, X_test, return_proba=return_proba)
         finally:
-            self.evict(dataset_id)
+            with self._id_locks.hold(dataset_id):
+                self._evict_from(dataset_id, self._locate(dataset_id), announce=False)
 
     def evict(self, dataset_id: str) -> None:
         with self._id_locks.hold(dataset_id):
-            for index in self._locate(dataset_id):
-                self._shards[index].evict(dataset_id)
-                self._emit("removed", dataset_id, index)
+            self._evict_from(dataset_id, self._locate(dataset_id), announce=True)
+
+    def residency(self) -> dict[str, dict[str, list[str]]]:
+        """Where every context is now, per device (lock-free snapshots):
+        what event subscribers reconcile against, since capacity eviction
+        without a spill tier drops contexts without an event."""
+        out = {}
+        for device, shard in zip(self._devices, self._shards, strict=True):
+            resident, known = shard.residency_nowait()
+            out[device] = {
+                "resident": sorted(resident),
+                "spilled": sorted(known - resident),
+            }
+        return out
 
     def stats(self) -> EngineStats:
         per_device = self.device_stats()
diff --git a/src/tabctx/serve/app.py b/src/tabctx/serve/app.py
index 59eef16..743af29 100644
--- a/src/tabctx/serve/app.py
+++ b/src/tabctx/serve/app.py
@@ -533,11 +533,17 @@ def models(self):
 
     @fastapi_app.get("/v1/tabctx/residency")
     def residency(self):
-        """Where contexts were placed and removed, latest last (multi-GPU
-        replicas): the router's decisions, for a platform scheduler or an
-        operator. Empty on a single-GPU replica, whose contexts all live
-        on its one device."""
-        return {"devices": list(self._devices), "events": self._residency.snapshot()}
+        """Where contexts are now, per device, and the router's latest
+        placements and removals (multi-GPU replicas): for a platform
+        scheduler or an operator. Events are advisory (contexts dropped by
+        capacity eviction without spill aren't reported); `devices` is the
+        current truth to reconcile against. A single-GPU replica reports
+        no per-device view or events."""
+        residency = getattr(self._engine, "residency", None)
+        return {
+            "devices": residency() if callable(residency) else {},
+            "events": self._residency.snapshot(),
+        }
 
     @fastapi_app.get("/v1/tabctx/limits")
     def limits(self):
diff --git a/src/tabctx/serve/workers.py b/src/tabctx/serve/workers.py
index a6278f7..219f947 100644
--- a/src/tabctx/serve/workers.py
+++ b/src/tabctx/serve/workers.py
@@ -509,6 +509,12 @@ def resolve_backend_name(self, backend: str | None) -> str:
     def holds(self, dataset_id: str) -> bool:
         return self.alive and dataset_id in self._snap().resident
 
+    def residency_nowait(self) -> tuple[frozenset[str], frozenset[str]]:
+        if not self.alive:
+            return frozenset(), frozenset()
+        snap = self._snap()
+        return snap.resident, snap.known
+
     def has(self, dataset_id: str) -> bool:
         return self.alive and dataset_id in self._snap().known
 
diff --git a/tests/unit/test_placement_events.py b/tests/unit/test_placement_events.py
index bd98d86..4a5eaad 100644
--- a/tests/unit/test_placement_events.py
+++ b/tests/unit/test_placement_events.py
@@ -106,3 +106,84 @@ def test_recent_events_are_bounded():
     for i in range(5):
         log(ResidencyEvent("placed", f"d{i}", "cuda:0"))
     assert [e["dataset_id"] for e in log.snapshot()] == ["d2", "d3", "d4"]
+
+
+def old_place(headrooms, needs, activity):
+    """The router's rule before policies were pluggable (607d3db)."""
+    indices = range(len(headrooms))
+    fits = [i for i in indices if headrooms[i] >= needs[i]]
+    if fits:
+        return min(fits, key=lambda i: (activity[i], -headrooms[i], i))
+    return max(indices, key=lambda i: (headrooms[i] - needs[i], -i))
+
+
+def test_default_policy_is_the_previous_rule_exactly():
+    import random
+
+    rng = random.Random(0)
+    for _ in range(5000):
+        n = rng.randint(2, 6)
+        headrooms = [rng.choice([0, 50, 100, 150, 400, 900]) for _ in range(n)]
+        needs = [rng.choice([1, 100, 200, 500, 1000]) for _ in range(n)]
+        activity = [(rng.randint(0, 2), rng.randint(0, 3)) for _ in range(n)]
+        views = [
+            view(i, headrooms[i], needs[i], activity[i][0], activity[i][1])
+            for i in range(n)
+        ]
+        assert LeastActiveThatFits().choose(REQ, views) == old_place(
+            headrooms, needs, activity
+        )
+
+
+def test_fit_predict_reports_nothing():
+    engine, _, _ = build(n_devices=2)
+    log = RecentEvents()
+    engine.events.subscribe(log)
+    engine.fit_predict(*table(10), [[1.0]])
+    assert log.snapshot() == []
+
+
+def test_a_failing_policy_fails_the_fit_without_an_event():
+    engine, shards, _ = build(n_devices=2)
+    log = RecentEvents()
+    engine.events.subscribe(log)
+
+    class Raises:
+        def choose(self, request, devices):
+            raise RuntimeError("scheduler unavailable")
+
+    engine._placement = Raises()
+    with pytest.raises(RuntimeError):
+        engine.fit(*table(10), dataset_id="d")
+    assert log.snapshot() == [] and where(shards, "d") == []
+
+
+def test_a_device_that_fails_to_evict_is_skipped_and_not_reported():
+    engine, shards, _ = build(n_devices=2)
+    log = RecentEvents()
+    engine.events.subscribe(log)
+    engine.fit(*table(10), dataset_id="d")
+    first = where(shards, "d")[0]
+    other = 1 - first
+    engine._placement = type("P", (), {"choose": lambda self, r, d: other})()
+
+    def dead(dataset_id):
+        raise RuntimeError("worker died")
+
+    shards[first].evict = dead
+    engine.fit(*table(10), dataset_id="d")  # the re-fit still succeeds
+    assert [e["kind"] for e in log.snapshot()] == ["placed", "placed"]
+
+
+def test_residency_reports_resident_and_spilled_per_device(tmp_path):
+    engine, shards, _ = build(n_devices=2, spill_root=tmp_path)
+    engine.fit(*table(10), dataset_id="a")
+    engine.fit(*table(10), dataset_id="b")
+    view_now = engine.residency()
+    assert set(view_now) == {"cuda:0", "cuda:1"}
+    resident = sorted(i for d in view_now.values() for i in d["resident"])
+    assert resident == ["a", "b"]
+    index = where(shards, "a")[0]
+    shards[index]._cache.evict_one(ignore_pins=True)  # capacity eviction: spills
+    after = engine.residency()[f"cuda:{index}"]
+    assert after["spilled"] and not set(after["spilled"]) & set(after["resident"])
diff --git a/tests/unit/test_yeo_johnson.py b/tests/unit/test_yeo_johnson.py
index 280d304..e3262f4 100644
--- a/tests/unit/test_yeo_johnson.py
+++ b/tests/unit/test_yeo_johnson.py
@@ -7,7 +7,18 @@
 pytest.importorskip("torch")
 from sklearn.preprocessing import PowerTransformer  # noqa: E402
 
-from tabctx.backends.yeo_johnson import FastPowerTransformer, fit_lambdas  # noqa: E402
+from tabctx.backends.yeo_johnson import (  # noqa: E402
+    FastPowerTransformer,
+    fit_lambdas,
+    on_device,
+)
+
+
+@pytest.fixture(autouse=True)
+def fast_path():
+    """The fast fit only runs inside on_device (CPU torch here)."""
+    with on_device(None):
+        yield
 
 
 def columns(n=4000, seed=0):
@@ -65,10 +76,104 @@ def test_lambdas_stay_within_scipys_overflow_bounds():
     assert np.all(np.isfinite(lam)) and np.all(np.isfinite(out))
 
 
-def test_the_fitted_transformer_pickles_as_a_power_transformer():
+def test_the_fitted_transformer_pickles_as_a_plain_power_transformer():
+    """Spilled contexts must not depend on this module's class path."""
     import pickle
 
-    fast = FastPowerTransformer(method="yeo-johnson").fit(columns())
-    again = pickle.loads(pickle.dumps(fast))
-    assert isinstance(again, PowerTransformer)
-    np.testing.assert_array_equal(again.lambdas_, fast.lambdas_)
+    X = columns()
+    fast = FastPowerTransformer(method="yeo-johnson").fit(X)
+    blob = pickle.dumps(fast)
+    assert b"tabctx" not in blob
+    again = pickle.loads(blob)
+    assert type(again) is PowerTransformer
+    np.testing.assert_array_equal(again.transform(X), fast.transform(X))
+
+
+def test_outside_on_device_it_is_scikit_learns(fast_path):
+    """A backend without fast_preprocessing shares the patched class."""
+    import threading
+
+    X = columns()
+    ref = PowerTransformer(method="yeo-johnson").fit(X)
+    box = {}
+    t = threading.Thread(  # a thread outside any on_device block
+        target=lambda: box.setdefault("l", FastPowerTransformer().fit(X).lambdas_)
+    )
+    t.start()
+    t.join()
+    np.testing.assert_array_equal(box["l"], ref.lambdas_)  # bit-identical: same code
+
+
+def test_the_device_section_runs_under_the_guard():
+    entered = []
+
+    class Guard:
+        def __enter__(self):
+            entered.append(True)
+
+        def __exit__(self, *exc):
+            return False
+
+    with on_device(None, Guard):
+        FastPowerTransformer(method="yeo-johnson").fit(columns())
+    assert entered == [True]
+
+
+def test_device_out_of_memory_falls_back_to_scikit_learn(monkeypatch):
+    import torch
+
+    import tabctx.backends.yeo_johnson as yj
+
+    def oom(*a, **k):
+        raise torch.cuda.OutOfMemoryError("CUDA out of memory")
+
+    monkeypatch.setattr(yj, "fit_lambdas", oom)
+    X = columns()
+    np.testing.assert_array_equal(
+        FastPowerTransformer(method="yeo-johnson").fit(X).lambdas_,
+        PowerTransformer(method="yeo-johnson").fit(X).lambdas_,
+    )
+
+
+def test_columns_are_fitted_in_bounded_chunks(monkeypatch):
+    import tabctx.backends.yeo_johnson as yj
+
+    X = columns()
+    whole = fit_lambdas(X[:, :-1])
+    monkeypatch.setattr(yj, "_CHUNK_BYTES", yj._BYTES_PER_CELL * X.shape[0] * 2)
+    calls = []
+    real = yj._fit_lambdas
+    monkeypatch.setattr(
+        yj, "_fit_lambdas", lambda x, d: calls.append(x.shape[1]) or real(x, d)
+    )
+    np.testing.assert_allclose(fit_lambdas(X[:, :-1]), whole, atol=1e-12)
+    assert calls == [2, 2, 2]  # 6 columns, 2 per chunk
+
+
+def test_values_beyond_any_safe_lambda_are_left_untransformed():
+    """Mixed-sign float32 values near 1e15: scipy raises; lambda 1 here."""
+    X = np.array([[1e15, 0.5], [-1e15, -0.5], [3e14, 0.1]] * 10, dtype=np.float32)
+    assert fit_lambdas(X)[0] == 1.0
+
+
+def test_float64_input_matches_scikit_learn():
+    X = columns().astype(np.float64)
+    np.testing.assert_allclose(
+        FastPowerTransformer(method="yeo-johnson").fit(X).lambdas_,
+        PowerTransformer(method="yeo-johnson").fit(X).lambdas_,
+        atol=1e-6,
+    )
+
+
+def test_install_replaces_tabicls_normalizer():
+    preprocessing = pytest.importorskip("tabicl._sklearn.preprocessing")
+    from tabctx.backends.yeo_johnson import install
+
+    original = preprocessing.PowerTransformer
+    try:
+        assert install()
+        pipeline = preprocessing.PreprocessingPipeline(normalization_method="power")
+        pipeline.fit(columns()[:, :-1])
+        assert isinstance(pipeline.normalizer_, FastPowerTransformer)
+    finally:
+        preprocessing.PowerTransformer = original

From d1d3607f9a2873828ed128ba183c58e78ecedd9e Mon Sep 17 00:00:00 2001
From: Amrit Krishnan <amrit110@gmail.com>
Date: Sun, 18 Oct 2026 01:24:43 -0400
Subject: [PATCH 34/39] bench: memory-pressure search with compiled blocks and
 fast preprocessing

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 benchmarks/bench_memory_pressure.py | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/benchmarks/bench_memory_pressure.py b/benchmarks/bench_memory_pressure.py
index 2ef76db..917c48e 100644
--- a/benchmarks/bench_memory_pressure.py
+++ b/benchmarks/bench_memory_pressure.py
@@ -106,7 +106,13 @@ def search(args, backend, rows: int, features: int) -> dict:
 def _measure_shape(args, rows: int, features: int, queue) -> None:
     """One shape in a fresh process: no allocator or library state left
     over from earlier shapes' OOMs can skew its measurements."""
-    backend = build_backend(args.backend, args.device, args.cache_mode)
+    backend = build_backend(
+        args.backend,
+        args.device,
+        args.cache_mode,
+        compile_blocks=args.compile,
+        fast_preprocessing=args.fast_preprocessing,
+    )
     backend.warmup()
     if args.search:
         queue.put([search(args, backend, rows, features)])
@@ -133,6 +139,10 @@ def main() -> int:
         help="with --search: a fit slower than this x unconstrained counts as "
         "failed (for models that degrade instead of failing, e.g. 2 for tabpfn)",
     )
+    p.add_argument("--compile", action="store_true", help="compiled TabICL blocks")
+    p.add_argument(
+        "--fast-preprocessing", action="store_true", help="vectorized Yeo-Johnson"
+    )
     p.add_argument("--out", default=None)
     args = p.parse_args()
 

From 025053c45160a9fb7eba61096525de9d11bfb07f Mon Sep 17 00:00:00 2001
From: Amrit Krishnan <amrit110@gmail.com>
Date: Sun, 18 Oct 2026 02:11:10 -0400
Subject: [PATCH 35/39] fix(memory): compiled TabICL is admitted against its
 own requirement grid

The minimum free memory a fit succeeds with (bench_memory_pressure
--search, A100) is unchanged or up to ~2 GiB higher with compiled blocks
(80k x 10: 22.9 vs 20.7 GiB; 50k x 100: 14.4 vs 13.0), although the
unconstrained peak is ~20% lower: tabicl sizes its batches from free
memory, so the floor, not the peak, is what admission reserves. Admitting
compiled fits on the eager grid could understate them near the limit.
The registry is now keyed by (model, cache mode, compiled); the compiled
TabICL grid takes, per shape, the larger of the eager and compiled
measurements. Corrects the roadmap's "so larger tables fit".

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 CHANGELOG.md                                  |   2 +-
 ROADMAP.md                                    |   5 +-
 .../minfree-tabicl-kv-compiled-a100.jsonl     |  16 +++
 .../minfree-tabicl-kv-compiled-max-a100.jsonl |  16 +++
 gke-out/minfree-tabicl-kv-compiled-a100.jsonl |  16 +++
 ...ration_tabicl_compiled_a100_requirement.py | 108 ++++++++++++++++++
 src/tabctx/serve/factory.py                   |  30 +++--
 tests/unit/test_serve_factory.py              |  28 ++++-
 8 files changed, 211 insertions(+), 10 deletions(-)
 create mode 100644 benchmarks/baselines/minfree-tabicl-kv-compiled-a100.jsonl
 create mode 100644 benchmarks/baselines/minfree-tabicl-kv-compiled-max-a100.jsonl
 create mode 100644 gke-out/minfree-tabicl-kv-compiled-a100.jsonl
 create mode 100644 src/tabctx/memory/calibration_tabicl_compiled_a100_requirement.py

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2688db6..31cfde6 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -16,7 +16,7 @@
 
 ### Changed
 - **TabICL's power transforms are fitted on the GPU** (`TABCTX_FAST_PREPROCESSING`, default on; `backends/yeo_johnson.py`). tabicl normalizes each ensemble member with scikit-learn's Yeo-Johnson PowerTransformer, which fits one lambda per column with scipy -- 80%+ of a fit's CPU half. A vectorized fit (same objective, bounds and tolerance; grid then golden-section search, all columns at once) runs in 0.08-0.25 s instead of 0.8-4.4 s on an A100 host; lambdas agree with scikit-learn's within one float32 ulp, accuracy/AUC/R^2 are unchanged on the evaluation datasets, and NaN inputs, Box-Cox and a device out-of-memory fall back to scikit-learn. The device section runs under the measurement lock (a concurrent fit's allocator readings never see it) with temporaries bounded to 512 MiB per column chunk, and fitted transformers pickle as plain scikit-learn objects. With compiled blocks, an isolated 40k x 100 fit takes 4.8 s (was 11.2 s).
-- **TabICL's attention blocks are compiled** (#36, `TABCTX_COMPILE`, default on; `backends/tabicl_compile.py`). Inductor fuses the memory-bound layer norms, residual adds and GELU MLP around each attention kernel: on an A100 the fit's GPU half is 1.7-2.2x faster (10k x 50: 0.91 -> 0.52 s; 40k x 100: 7.65 -> 4.30 s; 20k x 200: 6.51 -> 2.92 s) and its peak ~20% lower (40k x 100: 33.1 -> 27.0 GiB); 100-row predicts 20-35% faster. Accuracy, AUC and R^2 are unchanged on the evaluation datasets, and predicts beside a fit stay correct. Ray Serve, 6 mixed users on 4x A100: 5.78 -> 6.99 predicts/s, small-fit p95 1.48 -> 0.95 s. tabicl's ssmax log(n) is now computed symbolically inside compiled code, so a new table size doesn't recompile. Only tabicl's fp16 path (1,024+ rows) is compiled; its common graphs (~50) compile in `warmup()` at replica start (~130 s per worker, workers in parallel); shapes it doesn't cover (very large tables) compile on first use, ~10 s once per worker. Compile errors fall back to eager, with a warning after warmup if nothing compiled.
+- **TabICL's attention blocks are compiled** (#36, `TABCTX_COMPILE`, default on; `backends/tabicl_compile.py`). Inductor fuses the memory-bound layer norms, residual adds and GELU MLP around each attention kernel: on an A100 the fit's GPU half is 1.7-2.2x faster (10k x 50: 0.91 -> 0.52 s; 40k x 100: 7.65 -> 4.30 s; 20k x 200: 6.51 -> 2.92 s) and its peak ~20% lower (40k x 100: 33.1 -> 27.0 GiB); 100-row predicts 20-35% faster. Accuracy, AUC and R^2 are unchanged on the evaluation datasets, and predicts beside a fit stay correct. Ray Serve, 6 mixed users on 4x A100: 5.78 -> 6.99 predicts/s, small-fit p95 1.48 -> 0.95 s. tabicl's ssmax log(n) is now computed symbolically inside compiled code, so a new table size doesn't recompile. Only tabicl's fp16 path (1,024+ rows) is compiled; its common graphs (~50) compile in `warmup()` at replica start (~130 s per worker, workers in parallel); shapes it doesn't cover (very large tables) compile on first use, ~10 s once per worker. Compile errors fall back to eager, with a warning after warmup if nothing compiled. The lower peak doesn't lower what admission must reserve: the least free memory a compiled fit succeeds with is the same as eager's or up to ~2 GiB more (80k x 10: 22.9 vs 20.7 GiB), so compiled TabICL is admitted against its own requirement grid (per shape, the larger of the two measurements; `minfree-tabicl-kv-compiled-a100.jsonl`).
 - **Host memory is admitted by estimated bytes, node-wide** (`memory/host.py`). Each backend estimates a call's host use (TabICL ~5 KB per cell when it offloads stage outputs, else nothing; TabPFN ~0.6 KB per cell) and the guard admits calls against one budget shared by every worker process, first come first served (a large fit can't be starved by modest ones), with at most `slots` offloading calls at once (each budgets 1/slots of free host memory for staging). Calls under 1 GiB never wait. The old guard gave every call over 3M cells a whole 48 GiB slot, which serialized fits that use a few GiB (40k x 100 TabICL fits took p50 195 s under concurrent load on the 4x L4 rehearsal). A worker that dies holding host memory has it returned by its parent, and a stuck accounting lock fails calls with a retryable 503 instead of hanging. **`/readyz` `host_memory` keys changed**: `large_fits_active`/`large_fits_waiting` are now `calls_active`/`calls_waiting`, plus `reserved_bytes` and `calls_offloading`.
 - **TabICL fits prepare their CPU half before taking the device**: preprocessing (~30-37% of a fit) overlaps another fit's GPU half -- back-to-back fits 17-26% faster on an A100. At most one fit per device prepares ahead, so queued fits don't each hold preprocessed copies in host memory.
 - Worker processes shut down together under one 10 s deadline (was up to 10 s each, past a pod's grace period).
diff --git a/ROADMAP.md b/ROADMAP.md
index 77598ef..eb3f6dc 100644
--- a/ROADMAP.md
+++ b/ROADMAP.md
@@ -102,7 +102,10 @@ no FP8).
     forward compiled, so Inductor fuses the layer norms, residual adds and
     GELU MLP around the attention kernel. A100: GPU half 1.7-2.2x faster
     (40k x 100: 7.65 -> 4.30 s), peak memory ~20% lower (33.1 -> 27.0
-    GiB, so larger tables fit), accuracy unchanged on the evaluation sets;
+    GiB unconstrained -- but not the minimum a fit needs, which tabicl's
+    free-memory batch sizing sets; compiled needs up to ~2 GiB more at some
+    shapes, so it has its own requirement grid), accuracy unchanged on the
+    evaluation sets;
     Ray Serve, 6 users on 4x A100: 5.78 -> 6.99 predicts/s. Dynamic
     shapes needed tabicl's ssmax log(n) computed symbolically (math.log on
     the row count recompiled every new table). ~50 graphs (fp16 path
diff --git a/benchmarks/baselines/minfree-tabicl-kv-compiled-a100.jsonl b/benchmarks/baselines/minfree-tabicl-kv-compiled-a100.jsonl
new file mode 100644
index 0000000..864a2c4
--- /dev/null
+++ b/benchmarks/baselines/minfree-tabicl-kv-compiled-a100.jsonl
@@ -0,0 +1,16 @@
+{"rows": 2000, "features": 10, "status": "ok", "resident_gib": 0.39, "unconstrained_peak_gib": 0.59, "unconstrained_fit_s": 0.184, "min_free_gib": 0.69, "min_working_gib": 0.3, "fit_s_at_min": 0.199, "max_slowdown": null}
+{"rows": 20000, "features": 10, "status": "ok", "resident_gib": 3.68, "unconstrained_peak_gib": 5.53, "unconstrained_fit_s": 11.418, "min_free_gib": 5.6, "min_working_gib": 1.92, "fit_s_at_min": 0.956, "max_slowdown": null}
+{"rows": 45000, "features": 10, "status": "ok", "resident_gib": 8.27, "unconstrained_peak_gib": 12.41, "unconstrained_fit_s": 8.189, "min_free_gib": 12.55, "min_working_gib": 4.28, "fit_s_at_min": 3.269, "max_slowdown": null}
+{"rows": 80000, "features": 10, "status": "ok", "resident_gib": 14.68, "unconstrained_peak_gib": 22.04, "unconstrained_fit_s": 13.484, "min_free_gib": 22.94, "min_working_gib": 8.26, "fit_s_at_min": 8.573, "max_slowdown": null}
+{"rows": 2000, "features": 50, "status": "ok", "resident_gib": 0.45, "unconstrained_peak_gib": 1.09, "unconstrained_fit_s": 0.242, "min_free_gib": 1.05, "min_working_gib": 0.6, "fit_s_at_min": 0.265, "max_slowdown": null}
+{"rows": 10000, "features": 50, "status": "ok", "resident_gib": 1.94, "unconstrained_peak_gib": 5.14, "unconstrained_fit_s": 0.688, "min_free_gib": 3.09, "min_working_gib": 1.15, "fit_s_at_min": 0.718, "max_slowdown": null}
+{"rows": 40000, "features": 50, "status": "ok", "resident_gib": 7.4, "unconstrained_peak_gib": 17.6, "unconstrained_fit_s": 14.459, "min_free_gib": 11.8, "min_working_gib": 4.4, "fit_s_at_min": 11.828, "max_slowdown": null}
+{"rows": 2000, "features": 100, "status": "ok", "resident_gib": 0.52, "unconstrained_peak_gib": 1.93, "unconstrained_fit_s": 0.319, "min_free_gib": 1.12, "min_working_gib": 0.6, "fit_s_at_min": 1.017, "max_slowdown": null}
+{"rows": 20000, "features": 100, "status": "ok", "resident_gib": 3.81, "unconstrained_peak_gib": 16.08, "unconstrained_fit_s": 12.789, "min_free_gib": 5.99, "min_working_gib": 2.18, "fit_s_at_min": 8.869, "max_slowdown": null}
+{"rows": 50000, "features": 100, "status": "ok", "resident_gib": 9.33, "unconstrained_peak_gib": 24.02, "unconstrained_fit_s": 20.003, "min_free_gib": 14.38, "min_working_gib": 5.05, "fit_s_at_min": 29.394, "max_slowdown": null}
+{"rows": 1000, "features": 300, "status": "ok", "resident_gib": 0.63, "unconstrained_peak_gib": 2.86, "unconstrained_fit_s": 0.422, "min_free_gib": 1.23, "min_working_gib": 0.6, "fit_s_at_min": 1.691, "max_slowdown": null}
+{"rows": 5000, "features": 300, "status": "ok", "resident_gib": 1.38, "unconstrained_peak_gib": 11.43, "unconstrained_fit_s": 6.469, "min_free_gib": 2.26, "min_working_gib": 0.88, "fit_s_at_min": 7.077, "max_slowdown": null}
+{"rows": 1000, "features": 500, "status": "ok", "resident_gib": 0.92, "unconstrained_peak_gib": 4.68, "unconstrained_fit_s": 0.626, "min_free_gib": 1.51, "min_working_gib": 0.59, "fit_s_at_min": 2.475, "max_slowdown": null}
+{"rows": 5000, "features": 500, "status": "ok", "resident_gib": 1.68, "unconstrained_peak_gib": 16.45, "unconstrained_fit_s": 12.621, "min_free_gib": 2.84, "min_working_gib": 1.16, "fit_s_at_min": 12.255, "max_slowdown": null}
+{"rows": 500, "features": 1000, "status": "ok", "resident_gib": 1.56, "unconstrained_peak_gib": 5.36, "unconstrained_fit_s": 1.085, "min_free_gib": 2.43, "min_working_gib": 0.87, "fit_s_at_min": 2.469, "max_slowdown": null}
+{"rows": 2000, "features": 1000, "status": "ok", "resident_gib": 1.84, "unconstrained_peak_gib": 15.37, "unconstrained_fit_s": 12.801, "min_free_gib": 2.99, "min_working_gib": 1.15, "fit_s_at_min": 8.706, "max_slowdown": null}
diff --git a/benchmarks/baselines/minfree-tabicl-kv-compiled-max-a100.jsonl b/benchmarks/baselines/minfree-tabicl-kv-compiled-max-a100.jsonl
new file mode 100644
index 0000000..e29bf76
--- /dev/null
+++ b/benchmarks/baselines/minfree-tabicl-kv-compiled-max-a100.jsonl
@@ -0,0 +1,16 @@
+{"rows": 2000, "features": 10, "status": "ok", "resident_gib": 0.39, "unconstrained_peak_gib": 0.59, "unconstrained_fit_s": 0.433, "min_free_gib": 0.69, "min_working_gib": 0.3, "fit_s_at_min": 0.277}
+{"rows": 20000, "features": 10, "status": "ok", "resident_gib": 3.68, "unconstrained_peak_gib": 5.53, "unconstrained_fit_s": 11.418, "min_free_gib": 5.6, "min_working_gib": 1.92, "fit_s_at_min": 0.956, "max_slowdown": null}
+{"rows": 45000, "features": 10, "status": "ok", "resident_gib": 6.58, "unconstrained_peak_gib": 10.72, "unconstrained_fit_s": 4.581, "min_free_gib": 14.62, "min_working_gib": 8.04, "fit_s_at_min": 4.051}
+{"rows": 80000, "features": 10, "status": "ok", "resident_gib": 14.68, "unconstrained_peak_gib": 22.04, "unconstrained_fit_s": 13.484, "min_free_gib": 22.94, "min_working_gib": 8.26, "fit_s_at_min": 8.573, "max_slowdown": null}
+{"rows": 2000, "features": 50, "status": "ok", "resident_gib": 0.45, "unconstrained_peak_gib": 1.64, "unconstrained_fit_s": 0.694, "min_free_gib": 1.05, "min_working_gib": 0.6, "fit_s_at_min": 0.74}
+{"rows": 10000, "features": 50, "status": "ok", "resident_gib": 1.94, "unconstrained_peak_gib": 5.14, "unconstrained_fit_s": 0.688, "min_free_gib": 3.09, "min_working_gib": 1.15, "fit_s_at_min": 0.718, "max_slowdown": null}
+{"rows": 40000, "features": 50, "status": "ok", "resident_gib": 7.4, "unconstrained_peak_gib": 25.85, "unconstrained_fit_s": 6.619, "min_free_gib": 11.8, "min_working_gib": 4.4, "fit_s_at_min": 8.401}
+{"rows": 2000, "features": 100, "status": "ok", "resident_gib": 0.52, "unconstrained_peak_gib": 2.98, "unconstrained_fit_s": 1.273, "min_free_gib": 1.12, "min_working_gib": 0.6, "fit_s_at_min": 1.707}
+{"rows": 20000, "features": 100, "status": "ok", "resident_gib": 3.81, "unconstrained_peak_gib": 16.08, "unconstrained_fit_s": 12.789, "min_free_gib": 5.99, "min_working_gib": 2.18, "fit_s_at_min": 8.869, "max_slowdown": null}
+{"rows": 50000, "features": 100, "status": "ok", "resident_gib": 9.33, "unconstrained_peak_gib": 24.02, "unconstrained_fit_s": 20.003, "min_free_gib": 14.38, "min_working_gib": 5.05, "fit_s_at_min": 29.394, "max_slowdown": null}
+{"rows": 1000, "features": 300, "status": "ok", "resident_gib": 0.63, "unconstrained_peak_gib": 2.86, "unconstrained_fit_s": 0.422, "min_free_gib": 1.23, "min_working_gib": 0.6, "fit_s_at_min": 1.691, "max_slowdown": null}
+{"rows": 5000, "features": 300, "status": "ok", "resident_gib": 1.38, "unconstrained_peak_gib": 11.43, "unconstrained_fit_s": 6.469, "min_free_gib": 2.26, "min_working_gib": 0.88, "fit_s_at_min": 7.077, "max_slowdown": null}
+{"rows": 1000, "features": 500, "status": "ok", "resident_gib": 0.92, "unconstrained_peak_gib": 7.36, "unconstrained_fit_s": 5.105, "min_free_gib": 1.51, "min_working_gib": 0.59, "fit_s_at_min": 6.22}
+{"rows": 5000, "features": 500, "status": "ok", "resident_gib": 1.68, "unconstrained_peak_gib": 25.27, "unconstrained_fit_s": 9.848, "min_free_gib": 2.84, "min_working_gib": 1.16, "fit_s_at_min": 13.435}
+{"rows": 500, "features": 1000, "status": "ok", "resident_gib": 1.56, "unconstrained_peak_gib": 8.26, "unconstrained_fit_s": 9.29, "min_free_gib": 2.43, "min_working_gib": 0.87, "fit_s_at_min": 9.966}
+{"rows": 2000, "features": 1000, "status": "ok", "resident_gib": 1.84, "unconstrained_peak_gib": 24.56, "unconstrained_fit_s": 13.093, "min_free_gib": 2.99, "min_working_gib": 1.15, "fit_s_at_min": 16.456}
diff --git a/gke-out/minfree-tabicl-kv-compiled-a100.jsonl b/gke-out/minfree-tabicl-kv-compiled-a100.jsonl
new file mode 100644
index 0000000..864a2c4
--- /dev/null
+++ b/gke-out/minfree-tabicl-kv-compiled-a100.jsonl
@@ -0,0 +1,16 @@
+{"rows": 2000, "features": 10, "status": "ok", "resident_gib": 0.39, "unconstrained_peak_gib": 0.59, "unconstrained_fit_s": 0.184, "min_free_gib": 0.69, "min_working_gib": 0.3, "fit_s_at_min": 0.199, "max_slowdown": null}
+{"rows": 20000, "features": 10, "status": "ok", "resident_gib": 3.68, "unconstrained_peak_gib": 5.53, "unconstrained_fit_s": 11.418, "min_free_gib": 5.6, "min_working_gib": 1.92, "fit_s_at_min": 0.956, "max_slowdown": null}
+{"rows": 45000, "features": 10, "status": "ok", "resident_gib": 8.27, "unconstrained_peak_gib": 12.41, "unconstrained_fit_s": 8.189, "min_free_gib": 12.55, "min_working_gib": 4.28, "fit_s_at_min": 3.269, "max_slowdown": null}
+{"rows": 80000, "features": 10, "status": "ok", "resident_gib": 14.68, "unconstrained_peak_gib": 22.04, "unconstrained_fit_s": 13.484, "min_free_gib": 22.94, "min_working_gib": 8.26, "fit_s_at_min": 8.573, "max_slowdown": null}
+{"rows": 2000, "features": 50, "status": "ok", "resident_gib": 0.45, "unconstrained_peak_gib": 1.09, "unconstrained_fit_s": 0.242, "min_free_gib": 1.05, "min_working_gib": 0.6, "fit_s_at_min": 0.265, "max_slowdown": null}
+{"rows": 10000, "features": 50, "status": "ok", "resident_gib": 1.94, "unconstrained_peak_gib": 5.14, "unconstrained_fit_s": 0.688, "min_free_gib": 3.09, "min_working_gib": 1.15, "fit_s_at_min": 0.718, "max_slowdown": null}
+{"rows": 40000, "features": 50, "status": "ok", "resident_gib": 7.4, "unconstrained_peak_gib": 17.6, "unconstrained_fit_s": 14.459, "min_free_gib": 11.8, "min_working_gib": 4.4, "fit_s_at_min": 11.828, "max_slowdown": null}
+{"rows": 2000, "features": 100, "status": "ok", "resident_gib": 0.52, "unconstrained_peak_gib": 1.93, "unconstrained_fit_s": 0.319, "min_free_gib": 1.12, "min_working_gib": 0.6, "fit_s_at_min": 1.017, "max_slowdown": null}
+{"rows": 20000, "features": 100, "status": "ok", "resident_gib": 3.81, "unconstrained_peak_gib": 16.08, "unconstrained_fit_s": 12.789, "min_free_gib": 5.99, "min_working_gib": 2.18, "fit_s_at_min": 8.869, "max_slowdown": null}
+{"rows": 50000, "features": 100, "status": "ok", "resident_gib": 9.33, "unconstrained_peak_gib": 24.02, "unconstrained_fit_s": 20.003, "min_free_gib": 14.38, "min_working_gib": 5.05, "fit_s_at_min": 29.394, "max_slowdown": null}
+{"rows": 1000, "features": 300, "status": "ok", "resident_gib": 0.63, "unconstrained_peak_gib": 2.86, "unconstrained_fit_s": 0.422, "min_free_gib": 1.23, "min_working_gib": 0.6, "fit_s_at_min": 1.691, "max_slowdown": null}
+{"rows": 5000, "features": 300, "status": "ok", "resident_gib": 1.38, "unconstrained_peak_gib": 11.43, "unconstrained_fit_s": 6.469, "min_free_gib": 2.26, "min_working_gib": 0.88, "fit_s_at_min": 7.077, "max_slowdown": null}
+{"rows": 1000, "features": 500, "status": "ok", "resident_gib": 0.92, "unconstrained_peak_gib": 4.68, "unconstrained_fit_s": 0.626, "min_free_gib": 1.51, "min_working_gib": 0.59, "fit_s_at_min": 2.475, "max_slowdown": null}
+{"rows": 5000, "features": 500, "status": "ok", "resident_gib": 1.68, "unconstrained_peak_gib": 16.45, "unconstrained_fit_s": 12.621, "min_free_gib": 2.84, "min_working_gib": 1.16, "fit_s_at_min": 12.255, "max_slowdown": null}
+{"rows": 500, "features": 1000, "status": "ok", "resident_gib": 1.56, "unconstrained_peak_gib": 5.36, "unconstrained_fit_s": 1.085, "min_free_gib": 2.43, "min_working_gib": 0.87, "fit_s_at_min": 2.469, "max_slowdown": null}
+{"rows": 2000, "features": 1000, "status": "ok", "resident_gib": 1.84, "unconstrained_peak_gib": 15.37, "unconstrained_fit_s": 12.801, "min_free_gib": 2.99, "min_working_gib": 1.15, "fit_s_at_min": 8.706, "max_slowdown": null}
diff --git a/src/tabctx/memory/calibration_tabicl_compiled_a100_requirement.py b/src/tabctx/memory/calibration_tabicl_compiled_a100_requirement.py
new file mode 100644
index 0000000..bd14044
--- /dev/null
+++ b/src/tabctx/memory/calibration_tabicl_compiled_a100_requirement.py
@@ -0,0 +1,108 @@
+"""GENERATED by benchmarks/gen_requirement_grid.py -- do not hand-edit.
+
+Measured requirements (A100-SXM4-40GB, tabicl 2.2.0, kv, compiled attention blocks + fast preprocessing; per shape the larger of the eager and compiled requirement): for each shape, the least free GPU
+memory with which a fit succeeded (resident context + working floor), from
+benchmarks/bench_memory_pressure.py --search. See memory/requirement.py.
+"""
+
+from __future__ import annotations
+
+from tabctx.memory.adaptive import Observation
+from tabctx.memory.requirement import FloorPoint
+
+A100_40GB_TABICL_KV_COMPILED_FLOOR_POINTS: tuple[FloorPoint, ...] = (
+    FloorPoint(
+        n_train=2000, n_features=10, working_bytes=322122547, resident_bytes=418759311
+    ),  # resident 0.39 GiB, unconstrained peak 0.59 GiB, fit 0.433s -> 0.277s at the floor
+    FloorPoint(
+        n_train=20000,
+        n_features=10,
+        working_bytes=2061584302,
+        resident_bytes=3951369912,
+    ),  # resident 3.68 GiB, unconstrained peak 5.53 GiB, fit 11.418s -> 0.956s at the floor
+    FloorPoint(
+        n_train=45000,
+        n_features=10,
+        working_bytes=8632884264,
+        resident_bytes=7065221201,
+    ),  # resident 6.58 GiB, unconstrained peak 10.72 GiB, fit 4.581s -> 4.051s at the floor
+    FloorPoint(
+        n_train=80000,
+        n_features=10,
+        working_bytes=8869107466,
+        resident_bytes=15762529976,
+    ),  # resident 14.68 GiB, unconstrained peak 22.04 GiB, fit 13.484s -> 8.573s at the floor
+    FloorPoint(
+        n_train=2000, n_features=50, working_bytes=644245094, resident_bytes=483183820
+    ),  # resident 0.45 GiB, unconstrained peak 1.64 GiB, fit 0.694s -> 0.74s at the floor
+    FloorPoint(
+        n_train=10000,
+        n_features=50,
+        working_bytes=1234803097,
+        resident_bytes=2083059138,
+    ),  # resident 1.94 GiB, unconstrained peak 5.14 GiB, fit 0.688s -> 0.718s at the floor
+    FloorPoint(
+        n_train=40000,
+        n_features=50,
+        working_bytes=4724464025,
+        resident_bytes=7945689497,
+    ),  # resident 7.4 GiB, unconstrained peak 25.85 GiB, fit 6.619s -> 8.401s at the floor
+    FloorPoint(
+        n_train=2000, n_features=100, working_bytes=644245094, resident_bytes=558345748
+    ),  # resident 0.52 GiB, unconstrained peak 2.98 GiB, fit 1.273s -> 1.707s at the floor
+    FloorPoint(
+        n_train=20000,
+        n_features=100,
+        working_bytes=2340757176,
+        resident_bytes=4090956349,
+    ),  # resident 3.81 GiB, unconstrained peak 16.08 GiB, fit 12.789s -> 8.869s at the floor
+    FloorPoint(
+        n_train=50000,
+        n_features=100,
+        working_bytes=5422396211,
+        resident_bytes=10018011217,
+    ),  # resident 9.33 GiB, unconstrained peak 24.02 GiB, fit 20.003s -> 29.394s at the floor
+    FloorPoint(
+        n_train=1000, n_features=300, working_bytes=644245094, resident_bytes=676457349
+    ),  # resident 0.63 GiB, unconstrained peak 2.86 GiB, fit 0.422s -> 1.691s at the floor
+    FloorPoint(
+        n_train=5000, n_features=300, working_bytes=944892805, resident_bytes=1481763717
+    ),  # resident 1.38 GiB, unconstrained peak 11.43 GiB, fit 6.469s -> 7.077s at the floor
+    FloorPoint(
+        n_train=1000, n_features=500, working_bytes=633507676, resident_bytes=987842478
+    ),  # resident 0.92 GiB, unconstrained peak 7.36 GiB, fit 5.105s -> 6.22s at the floor
+    FloorPoint(
+        n_train=5000,
+        n_features=500,
+        working_bytes=1245540515,
+        resident_bytes=1803886264,
+    ),  # resident 1.68 GiB, unconstrained peak 25.27 GiB, fit 9.848s -> 13.435s at the floor
+    FloorPoint(
+        n_train=500, n_features=1000, working_bytes=934155386, resident_bytes=1675037245
+    ),  # resident 1.56 GiB, unconstrained peak 8.26 GiB, fit 9.29s -> 9.966s at the floor
+    FloorPoint(
+        n_train=2000,
+        n_features=1000,
+        working_bytes=1234803097,
+        resident_bytes=1975684956,
+    ),  # resident 1.84 GiB, unconstrained peak 24.56 GiB, fit 13.093s -> 16.456s at the floor
+)
+
+A100_40GB_TABICL_KV_COMPILED_REQUIREMENT_GRID: tuple[Observation, ...] = (
+    Observation(n_train=2000, n_features=10, real_bytes=740881858),
+    Observation(n_train=20000, n_features=10, real_bytes=6012954214),
+    Observation(n_train=45000, n_features=10, real_bytes=15698105466),
+    Observation(n_train=80000, n_features=10, real_bytes=24631637442),
+    Observation(n_train=2000, n_features=50, real_bytes=1127428915),
+    Observation(n_train=10000, n_features=50, real_bytes=3317862236),
+    Observation(n_train=40000, n_features=50, real_bytes=12670153523),
+    Observation(n_train=2000, n_features=100, real_bytes=1202590842),
+    Observation(n_train=20000, n_features=100, real_bytes=6431713525),
+    Observation(n_train=50000, n_features=100, real_bytes=15440407429),
+    Observation(n_train=1000, n_features=300, real_bytes=1320702443),
+    Observation(n_train=5000, n_features=300, real_bytes=2426656522),
+    Observation(n_train=1000, n_features=500, real_bytes=1621350154),
+    Observation(n_train=5000, n_features=500, real_bytes=3049426780),
+    Observation(n_train=500, n_features=1000, real_bytes=2609192632),
+    Observation(n_train=2000, n_features=1000, real_bytes=3210488053),
+)
diff --git a/src/tabctx/serve/factory.py b/src/tabctx/serve/factory.py
index 4b15062..4e36dbe 100644
--- a/src/tabctx/serve/factory.py
+++ b/src/tabctx/serve/factory.py
@@ -337,11 +337,26 @@ def _preloaded_observations(
 
 
 # Generated requirement calibrations (benchmarks/gen_requirement_grid.py),
-# per (model, cache mode): module under tabctx.memory and its name prefix.
-_REQUIREMENT_CALIBRATIONS: dict[tuple[str, str], tuple[str, str]] = {
-    ("tabicl", "kv"): ("calibration_tabicl_a100_requirement", "A100_40GB_TABICL_KV"),
-    ("tabpfn", "kv"): ("calibration_tabpfn_a100_requirement", "A100_40GB_TABPFN_KV"),
+# per (model, cache mode, compiled): module under tabctx.memory and its name
+# prefix. Compiled TabICL needs up to ~2 GiB more than eager at some shapes
+# (its minimum, not its peak: tabicl sizes batches from free memory), so it
+# has its own grid -- per shape the larger of the two measurements.
+_REQUIREMENT_CALIBRATIONS: dict[tuple[str, str, bool], tuple[str, str]] = {
+    ("tabicl", "kv", False): (
+        "calibration_tabicl_a100_requirement",
+        "A100_40GB_TABICL_KV",
+    ),
+    ("tabicl", "kv", True): (
+        "calibration_tabicl_compiled_a100_requirement",
+        "A100_40GB_TABICL_KV_COMPILED",
+    ),
+    ("tabpfn", "kv", False): (
+        "calibration_tabpfn_a100_requirement",
+        "A100_40GB_TABPFN_KV",
+    ),
 }
+# Backends whose compile_models setting changes their memory profile.
+_COMPILED_KINDS = frozenset({"tabicl"})
 # Models whose per-row resident cache grows with features (see
 # ResidentRateModel): TabPFN caches per (row, feature group); TabICL's ICL
 # cache is per row only.
@@ -349,12 +364,13 @@ def _preloaded_observations(
 
 
 def _requirement_calibration(
-    kind: BackendKind, kv_cache: KvCacheMode
+    kind: BackendKind, kv_cache: KvCacheMode, compiled: bool = False
 ) -> tuple[WorkingFloorModel | None, tuple, tuple]:
     """Measured working floors + requirements (memory/requirement.py), so
     admission bounds resident + floor instead of an unconstrained peak.
     Configurations without a measured calibration keep peak semantics."""
-    entry = _REQUIREMENT_CALIBRATIONS.get((kind, kv_cache))
+    compiled = compiled and kind in _COMPILED_KINDS
+    entry = _REQUIREMENT_CALIBRATIONS.get((kind, kv_cache, compiled))
     if entry is None:
         return None, (), ()
     module_name, prefix = entry
@@ -471,7 +487,7 @@ def build_estimator(
     )
     fit_grid, predict_grid = _preloaded_observations(kind, settings.kv_cache)
     working_floor, requirements, floor_points = _requirement_calibration(
-        kind, settings.kv_cache
+        kind, settings.kv_cache, settings.compile_models
     )
     return AdaptiveMemoryEstimator(
         fallback=fallback,
diff --git a/tests/unit/test_serve_factory.py b/tests/unit/test_serve_factory.py
index b39e047..e69384f 100644
--- a/tests/unit/test_serve_factory.py
+++ b/tests/unit/test_serve_factory.py
@@ -324,7 +324,11 @@ def test_missing_calibration_module_degrades_to_no_preload(self, monkeypatch):
         # shortcut finds the real module other tests already imported.
         import tabctx.memory as memory_pkg
 
-        for name in ("calibration_tabicl_a100", "calibration_tabicl_a100_requirement"):
+        for name in (
+            "calibration_tabicl_a100",
+            "calibration_tabicl_a100_requirement",
+            "calibration_tabicl_compiled_a100_requirement",
+        ):
             monkeypatch.delattr(memory_pkg, name, raising=False)
             monkeypatch.setitem(sys.modules, f"tabctx.memory.{name}", None)
         est = build_estimator(ServeSettings(backends=("tabicl",)))
@@ -490,3 +494,25 @@ def test_fast_preprocessing_flag_from_env(monkeypatch):
     assert ServeSettings.from_env().fast_preprocessing
     monkeypatch.setenv(FAST_PREPROCESSING_ENV_VAR, "0")
     assert not ServeSettings.from_env().fast_preprocessing
+
+
+def test_compiled_tabicl_uses_its_own_requirement_grid():
+    """Compiled TabICL's minimum need exceeds eager's at some shapes (80k x
+    10: 22.9 vs 20.7 GiB on an A100), so it must not inherit eager's grid."""
+    from tabctx.serve.factory import _requirement_calibration
+
+    _, eager, _ = _requirement_calibration("tabicl", "kv", compiled=False)
+    _, compiled, _ = _requirement_calibration("tabicl", "kv", compiled=True)
+
+    def need(grid, n, d):
+        return next(o.real_bytes for o in grid if (o.n_train, o.n_features) == (n, d))
+
+    assert need(compiled, 80_000, 10) > need(eager, 80_000, 10)
+    # ...and never less than eager anywhere (per shape, the larger of both).
+    for o in eager:
+        assert need(compiled, o.n_train, o.n_features) >= o.real_bytes
+    # TabPFN isn't compiled: the flag doesn't change its grid.
+    assert (
+        _requirement_calibration("tabpfn", "kv", True)[1]
+        == (_requirement_calibration("tabpfn", "kv", False)[1])
+    )

From 4e7a9508c77c554663f5c414f8d4a45b1f384ed4 Mon Sep 17 00:00:00 2001
From: Amrit Krishnan <amrit110@gmail.com>
Date: Sun, 18 Oct 2026 02:20:07 -0400
Subject: [PATCH 36/39] bench+docs: predict throughput is host-bound (~21/s per
 GPU, GPU 20% busy); roadmap status for #33, #35, #37

benchmarks/bench_predict_throughput.py measures predicts/s, latency and GPU
utilization with T users against distinct cached contexts on one device:
the measurement that redirects #33 from cross-context GPU batching to
host overhead (CUDA graphs for predict, or more processes per GPU).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 ROADMAP.md                             |  26 +++-
 benchmarks/bench_predict_throughput.py | 170 +++++++++++++++++++++++++
 2 files changed, 190 insertions(+), 6 deletions(-)
 create mode 100644 benchmarks/bench_predict_throughput.py

diff --git a/ROADMAP.md b/ROADMAP.md
index eb3f6dc..7bf5ef4 100644
--- a/ROADMAP.md
+++ b/ROADMAP.md
@@ -82,7 +82,9 @@ no FP8).
    own requirement grid (A100, 20 shapes): it, too, sizes working memory
    from what is free, so admission reaches 100k rows at 100 features on an
    A40 (was 5k on an L4 under TabICL's formula). Still to do: L4/A40-
-   measured grids (A100 grids scale by capacity today).
+   measured grids (A100 grids scale by capacity today; A40 time on GKE doesn't exist, so
+   the A40 grid comes from prod once tabctx runs there). Compiled TabICL
+   has its own A100 grid: its minimum need is up to ~2 GiB above eager's.
 8. **Head-of-line scheduling** ([#39](https://github.com/VectorInstitute/tabctx/issues/39), done): light predicts (resident
    context, <= 2048 rows) run beside a fit on the same GPU instead of
    waiting for it -- ~0.1 s p50 during a 40k x 100 fit instead of its
@@ -91,13 +93,25 @@ no FP8).
    wait. Small-fit p95 under load 10.6 s -> 1.4 s.
 9. **Tiled/fused column+row stage kernels** ([#32](https://github.com/VectorInstitute/tabctx/issues/32)): cut the fit peak and the
    memory-bound layer norms, raising the per-GPU table ceiling.
-10. **Paged context KV + cross-context varlen predict** ([#33](https://github.com/VectorInstitute/tabctx/issues/33)): throughput
-    with many users (PagedAttention + continuous batching analogue).
+10. **Paged context KV + cross-context varlen predict** ([#33](https://github.com/VectorInstitute/tabctx/issues/33), measured;
+    redirected): batching several tenants' predicts into one GPU call only
+    pays if predicts are GPU-bound, and they aren't. A100, compiled, 32
+    contexts of 2,000 x 20, 100-row predicts (`benchmarks/
+    bench_predict_throughput.py`): ~21 predicts/s per GPU at 1 to 8
+    concurrent users, latency growing linearly, GPU ~20% busy. A predict
+    is ~47 ms of host-side kernel launches (dynamo dispatch, Triton and
+    cuBLAS launches over ~36 block calls), so the levers are host
+    overhead: CUDA graphs for the predict path (shapes vary: bucket test
+    rows), or two worker processes per GPU for predict-heavy loads.
+    Cross-context KV paging matters again once predicts are launch-free.
 11. **INT8 context KV** ([#34](https://github.com/VectorInstitute/tabctx/issues/34), done for TabICL): ~0.52x resident, identical
     labels on the evaluation sets; TabPFN already defaults to int8.
-12. **Context parallelism across GPUs** ([#35](https://github.com/VectorInstitute/tabctx/issues/35)): tables larger than one GPU
-    (ring attention over the replica's devices). Disk offload already
-    lets one A100 fit 130k x 500 (65M cells) in 26 min.
+12. **Context parallelism across GPUs** ([#35](https://github.com/VectorInstitute/tabctx/issues/35), not started): tables larger
+    than one GPU (ring attention over the replica's devices). Needs model-
+    level changes in tabicl's attention (its InferenceManager owns
+    batching and offload), so it waits for a workload that exceeds one
+    A40: disk offload already lets one A100 fit 130k x 500 (65M cells) in
+    26 min, and prod's A40s have 48 GB.
 13. **torch.compile** ([#36](https://github.com/VectorInstitute/tabctx/issues/36), done for TabICL): each attention block's
     forward compiled, so Inductor fuses the layer norms, residual adds and
     GELU MLP around the attention kernel. A100: GPU half 1.7-2.2x faster
diff --git a/benchmarks/bench_predict_throughput.py b/benchmarks/bench_predict_throughput.py
new file mode 100644
index 0000000..2a99ef6
--- /dev/null
+++ b/benchmarks/bench_predict_throughput.py
@@ -0,0 +1,170 @@
+#!/usr/bin/env python3
+"""Predict throughput against many cached contexts on one GPU (#33).
+
+Before building cross-context batched predict (one kernel launch serving
+several tenants' predicts, the continuous-batching analogue), measure
+whether there is anything to win: fit N contexts, then T threads each
+predict R rows against distinct contexts, through one TabctxEngine (whose
+predict lock runs one GPU predict at a time). Reports predicts/s and
+latency per T, and the GPU's utilization while it runs (nvidia-smi).
+If throughput stays flat with T while utilization is low, predicts are
+launch/host bound and batching across contexts would pay; if the GPU is
+already busy, it wouldn't.
+
+Usage (on a GPU host):
+  bench_predict_throughput.py --contexts 32 --shape 2000x20 --rows 100 \
+      --threads 1 2 4 8 --seconds 20 --compile --fast-preprocessing
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import subprocess
+import sys
+import threading
+import time
+
+import numpy as np
+
+sys.path.insert(0, __file__.rsplit("/", 1)[0])
+from profile_fit_predict import build_backend, make_table, parse_shape  # noqa: E402
+
+from tabctx.cache.manager import ContextCacheManager  # noqa: E402
+from tabctx.engine import TabctxEngine  # noqa: E402
+from tabctx.memory import (  # noqa: E402
+    A100_40GB_TABICL_CALIBRATION,
+    AdaptiveMemoryEstimator,
+    PowerLawMemoryEstimator,
+)
+
+
+class Utilization:
+    """Samples GPU utilization with nvidia-smi while running."""
+
+    def __init__(self, index: int) -> None:
+        self._index = index
+        self._samples: list[int] = []
+        self._stop = threading.Event()
+        self._thread = threading.Thread(target=self._run, daemon=True)
+
+    def _run(self) -> None:
+        while not self._stop.wait(0.2):
+            out = subprocess.run(
+                [
+                    "nvidia-smi",
+                    f"--id={self._index}",
+                    "--query-gpu=utilization.gpu",
+                    "--format=csv,noheader,nounits",
+                ],
+                capture_output=True,
+                text=True,
+            )
+            if out.returncode == 0 and out.stdout.strip().isdigit():
+                self._samples.append(int(out.stdout.strip()))
+
+    def __enter__(self):
+        self._thread.start()
+        return self
+
+    def __exit__(self, *exc):
+        self._stop.set()
+        self._thread.join(5)
+
+    @property
+    def mean(self) -> float:
+        return round(float(np.mean(self._samples)), 1) if self._samples else 0.0
+
+
+def run_level(engine, ids, X_test, threads: int, seconds: float, index: int):
+    """T threads predicting against distinct contexts for `seconds`:
+    (latencies, elapsed, mean GPU utilization)."""
+    latencies: list[float] = []
+    lock = threading.Lock()
+    deadline = time.perf_counter() + seconds
+
+    def user(k: int) -> None:
+        n = 0
+        while time.perf_counter() < deadline:
+            dataset_id = ids[(k + n * threads) % len(ids)]
+            t0 = time.perf_counter()
+            engine.predict(dataset_id, X_test)
+            with lock:
+                latencies.append(time.perf_counter() - t0)
+            n += 1
+
+    with Utilization(index) as util:
+        start = time.perf_counter()
+        workers = [threading.Thread(target=user, args=(k,)) for k in range(threads)]
+        for w in workers:
+            w.start()
+        for w in workers:
+            w.join()
+        elapsed = time.perf_counter() - start
+    return np.array(latencies), elapsed, util.mean
+
+
+def main() -> int:
+    p = argparse.ArgumentParser()
+    p.add_argument("--backend", choices=["tabicl", "tabpfn"], default="tabicl")
+    p.add_argument("--device", default="cuda:0")
+    p.add_argument("--contexts", type=int, default=32)
+    p.add_argument("--shape", type=parse_shape, default=(2000, 20))
+    p.add_argument("--rows", type=int, default=100)
+    p.add_argument("--threads", type=int, nargs="+", default=[1, 2, 4, 8])
+    p.add_argument("--seconds", type=float, default=20.0)
+    p.add_argument("--compile", action="store_true")
+    p.add_argument("--fast-preprocessing", action="store_true")
+    p.add_argument("--out")
+    args = p.parse_args()
+
+    backend = build_backend(
+        args.backend,
+        args.device,
+        "kv",
+        compile_blocks=args.compile,
+        fast_preprocessing=args.fast_preprocessing,
+    )
+    backend.warmup()
+    estimator = AdaptiveMemoryEstimator(
+        fallback=PowerLawMemoryEstimator(A100_40GB_TABICL_CALIBRATION)
+    )
+    engine = TabctxEngine(
+        backend=backend,
+        cache=ContextCacheManager(estimator.ceiling_bytes()),
+        estimator=estimator,
+    )
+    rows, features = args.shape
+    ids = []
+    for i in range(args.contexts):
+        X, y = make_table(rows, features, "classification", seed=i)
+        ids.append(engine.fit(X, y, "classification", dataset_id=f"ctx{i}"))
+    X_test, _ = make_table(args.rows, features, "classification", seed=999)
+    for dataset_id in ids:  # warm every context's predict path
+        engine.predict(dataset_id, X_test)
+
+    index = int(args.device.split(":")[1]) if ":" in args.device else 0
+    for threads in args.threads:
+        lat, elapsed, util = run_level(
+            engine, ids, X_test, threads, args.seconds, index
+        )
+        record = {
+            "backend": args.backend,
+            "compiled": args.compile,
+            "shape": f"{rows}x{features}",
+            "rows": args.rows,
+            "threads": threads,
+            "predicts_per_s": round(len(lat) / elapsed, 2),
+            "p50_ms": round(float(np.percentile(lat, 50)) * 1000, 1),
+            "p95_ms": round(float(np.percentile(lat, 95)) * 1000, 1),
+            "gpu_util_pct": util,
+        }
+        print(json.dumps(record), flush=True)
+        if args.out:
+            with open(args.out, "a") as f:
+                f.write(json.dumps(record) + "\n")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

From 454a85aea3bf160894e518cf687848ba2232b038 Mon Sep 17 00:00:00 2001
From: Amrit Krishnan <amrit110@gmail.com>
Date: Sun, 18 Oct 2026 02:28:06 -0400
Subject: [PATCH 37/39] fix: final pre-release review of 0.12.0

- Host budget is re-measured once per-GPU workers have warmed up, and
  subtracts the container's working set (every process in the pod: the
  workers' models, CUDA contexts and compiled kernels) rather than only the
  serving process's RSS; the budget and slots live in the guard's shared
  state so every worker sees the update.
- Each worker sizes its torch.compile pool to its share of the container's
  CPUs (TORCHINDUCTOR_COMPILE_THREADS when unset): four workers compiling at
  startup oversubscribed the pod.
- /v1/tabctx/residency is tenant-scoped like every /v1/tabctx route: a
  caller sees only its own datasets (it listed every tenant's ids).
- The legacy predict endpoint no longer creates a CUDA context per GPU in
  the serving process when worker processes run the call.
- Cache docstring no longer describes the removed replica-wide lock; the
  unused ContextCacheManager.lock is gone; gke-out/ removed and ignored.
- CHANGELOG: int8 TabICL KV by default (served probabilities change within
  0.009; native for exact), TABCTX_OFFLOAD_DIR, BuiltEngine changes; the
  host guard described as new (it never shipped before), not as a change.
  README: int8 cache and offload dir. ROADMAP: Dynamo wording.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .gitignore                                    |   1 +
 CHANGELOG.md                                  |   9 +-
 README.md                                     |   9 ++
 ROADMAP.md                                    |   4 +-
 gke-out/minfree-tabicl-kv-compiled-a100.jsonl |  16 ---
 src/tabctx/cache/manager.py                   |  19 +---
 src/tabctx/memory/host.py                     | 101 ++++++++++++++----
 src/tabctx/serve/app.py                       |  42 ++++++--
 src/tabctx/serve/factory.py                   |   6 ++
 src/tabctx/serve/tenancy.py                   |  10 ++
 src/tabctx/serve/workers.py                   |  28 +++++
 tests/unit/test_host_memory.py                |  40 +++++++
 tests/unit/test_serve_tenancy.py              |  12 +++
 tests/unit/test_workers.py                    |  14 +++
 14 files changed, 247 insertions(+), 64 deletions(-)
 delete mode 100644 gke-out/minfree-tabicl-kv-compiled-a100.jsonl

diff --git a/.gitignore b/.gitignore
index 2339f31..d6b20e4 100644
--- a/.gitignore
+++ b/.gitignore
@@ -7,3 +7,4 @@ build/
 .pytest_cache/
 .env
 .coverage
+gke-out/
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 31cfde6..a33da89 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,16 +8,19 @@
 - **One worker process per GPU** (#39, `TABCTX_WORKER_PROCESSES`, default on for multi-GPU replicas). Fits are host-bound, so threads of one process driving several GPUs contend for the GIL: on 4x A100, four concurrent fits ran at 2.1-3.1x aggregate from threads and 3.96-4.3x from processes (`benchmarks/bench_scaling.py`). `serve/workers.py` runs each device's engine in its own process behind a `ProcessShard` proxy; the router reads versioned residency/headroom snapshots that every reply carries, so its decisions stay lock-free. A dead worker fails its calls, is skipped by placement, and fails `/readyz` and Ray Serve's `check_health` so the replica restarts. Nothing depends on Ray. Under Ray Serve on 4x A100 with 6 mixed users: 2.58 -> 5.0 predicts/s, small-fit p50 1.6 -> 0.6 s.
 - **Light predicts run beside a fit** on the same GPU (#39). A predict against a resident context with at most 2,048 rows no longer waits for a fit on its device (it waited up to 266 s on the L4 rehearsal): ~0.1 s p50 during a 40k x 100 fit on an A100, at ~15% fit slowdown. Restores and large predicts still wait; contexts in use are pinned against evict-ahead (and evicted with predicts held off when nothing else makes room); a fit that runs out of memory beside a predict retries alone; allocator measurements are taken with predicts quiesced, and memory freed during a fit is released only after it has measured itself.
 - **TabPFN-3.5 requirement calibration** (A100, 20 shapes up to 200k x 10 and 30k x 1,000). TabPFN also sizes its working memory from what is free (20k x 100 peaks at 7.3 GiB unconstrained but runs within 2x its time on 1.1 GiB), so admission bounds resident + working floor for it too, instead of TabICL's peak formula: up to 100k rows at 100 features and 80k at 1,000 on an A40 (was 5k and 500 on an L4). Its per-row resident rate grows with features, which `ResidentRateModel(width_dependent=True)` respects.
-- **Pluggable placement and residency events** (#39). The multi-device router asks a `PlacementPolicy` (`placement.py`) which device takes a fit, from a lock-free view of each device (headroom, the fit's estimated need there, activity, liveness); the previous rule is the default `LeastActiveThatFits`, which now also skips dead workers. The router emits `ResidencyEvent`s (`events.py`) when it places a context and when it removes a copy, to any sink; `/v1/tabctx/residency` serves the latest. A platform that co-schedules LLMs and tabular models can supply its own policy and follow placements without polling; events are advisory (contexts dropped by capacity eviction without spill aren't reported), so `/v1/tabctx/residency` also returns each device's current resident and spilled ids to reconcile against. `fit_predict`'s throwaway contexts aren't reported, and an evict that fails on one device (a dead worker) no longer stops the others.
+- **Pluggable placement and residency events** (#39). The multi-device router asks a `PlacementPolicy` (`placement.py`) which device takes a fit, from a lock-free view of each device (headroom, the fit's estimated need there, activity, liveness); the previous rule is the default `LeastActiveThatFits`, which now also skips dead workers. The router emits `ResidencyEvent`s (`events.py`) when it places a context and when it removes a copy, to any sink; `/v1/tabctx/residency` serves the latest. A platform that co-schedules LLMs and tabular models can supply its own policy and follow placements without polling; events are advisory (contexts dropped by capacity eviction without spill aren't reported), so `/v1/tabctx/residency` also returns each device's current resident and spilled ids to reconcile against, scoped like every `/v1/tabctx` route (a tenant sees only its own datasets). `fit_predict`'s throwaway contexts aren't reported, and an evict that fails on one device (a dead worker) no longer stops the others.
 - `benchmarks/bench_scaling.py` (threads vs processes across GPUs), `bench_overlap.py` (predict latency and correctness during a fit), `gke/serve_workload.sh` (Ray Serve + workload inside the devbox).
 - `benchmarks/bench_workload.py`: concurrent users looping fit-then-predict over a small/medium/large table mix (large ones via CSV upload), with per-class latency percentiles, 413s, errors and the per-device spread.
 - `benchmarks/sweep_admission.py`: the table-size sweep (rows up to the first 413 per feature count, fit and predict timed).
 - ROADMAP: multi-GPU and CUDA-level track (#30-#37), PyTorch-level first.
 
 ### Changed
+- **TabICL's ICL key/value cache is stored in int8 by default** (#34, `TABCTX_KV_PRECISION=int8|fp16|native`; `backends/tabicl_kv.py`). It dominates a TabICL context, so int8 about halves resident memory per context: twice the contexts per GPU and a longer longest table. A100, 7 datasets (`benchmarks/baselines/kv-precision-tabicl-a100.jsonl`): identical labels and metrics, max probability drift 0.009, warm predicts up to ~1.4x slower on large contexts. **Served probabilities change slightly**; set `TABCTX_KV_PRECISION=native` for 0.11.0's exact outputs. TabPFN already stores its cache in int8.
+- **`TABCTX_OFFLOAD_DIR`** (new): where TabICL spills stage outputs too large for GPU and host memory (memory-mapped files) instead of exhausting the pod's RAM; default `<TABCTX_SPILL_DIR>/offload`, else a temporary directory, one subdirectory per device. Lets one A100 fit 130k x 500 (65M cells).
+- `serve.factory.BuiltEngine`: `backends`/`backend` are replaced by `backend_names` (with worker processes this process holds no backends), `engine` may be a `MultiDeviceEngine`, and `close()` stops worker processes.
 - **TabICL's power transforms are fitted on the GPU** (`TABCTX_FAST_PREPROCESSING`, default on; `backends/yeo_johnson.py`). tabicl normalizes each ensemble member with scikit-learn's Yeo-Johnson PowerTransformer, which fits one lambda per column with scipy -- 80%+ of a fit's CPU half. A vectorized fit (same objective, bounds and tolerance; grid then golden-section search, all columns at once) runs in 0.08-0.25 s instead of 0.8-4.4 s on an A100 host; lambdas agree with scikit-learn's within one float32 ulp, accuracy/AUC/R^2 are unchanged on the evaluation datasets, and NaN inputs, Box-Cox and a device out-of-memory fall back to scikit-learn. The device section runs under the measurement lock (a concurrent fit's allocator readings never see it) with temporaries bounded to 512 MiB per column chunk, and fitted transformers pickle as plain scikit-learn objects. With compiled blocks, an isolated 40k x 100 fit takes 4.8 s (was 11.2 s).
-- **TabICL's attention blocks are compiled** (#36, `TABCTX_COMPILE`, default on; `backends/tabicl_compile.py`). Inductor fuses the memory-bound layer norms, residual adds and GELU MLP around each attention kernel: on an A100 the fit's GPU half is 1.7-2.2x faster (10k x 50: 0.91 -> 0.52 s; 40k x 100: 7.65 -> 4.30 s; 20k x 200: 6.51 -> 2.92 s) and its peak ~20% lower (40k x 100: 33.1 -> 27.0 GiB); 100-row predicts 20-35% faster. Accuracy, AUC and R^2 are unchanged on the evaluation datasets, and predicts beside a fit stay correct. Ray Serve, 6 mixed users on 4x A100: 5.78 -> 6.99 predicts/s, small-fit p95 1.48 -> 0.95 s. tabicl's ssmax log(n) is now computed symbolically inside compiled code, so a new table size doesn't recompile. Only tabicl's fp16 path (1,024+ rows) is compiled; its common graphs (~50) compile in `warmup()` at replica start (~130 s per worker, workers in parallel); shapes it doesn't cover (very large tables) compile on first use, ~10 s once per worker. Compile errors fall back to eager, with a warning after warmup if nothing compiled. The lower peak doesn't lower what admission must reserve: the least free memory a compiled fit succeeds with is the same as eager's or up to ~2 GiB more (80k x 10: 22.9 vs 20.7 GiB), so compiled TabICL is admitted against its own requirement grid (per shape, the larger of the two measurements; `minfree-tabicl-kv-compiled-a100.jsonl`).
-- **Host memory is admitted by estimated bytes, node-wide** (`memory/host.py`). Each backend estimates a call's host use (TabICL ~5 KB per cell when it offloads stage outputs, else nothing; TabPFN ~0.6 KB per cell) and the guard admits calls against one budget shared by every worker process, first come first served (a large fit can't be starved by modest ones), with at most `slots` offloading calls at once (each budgets 1/slots of free host memory for staging). Calls under 1 GiB never wait. The old guard gave every call over 3M cells a whole 48 GiB slot, which serialized fits that use a few GiB (40k x 100 TabICL fits took p50 195 s under concurrent load on the 4x L4 rehearsal). A worker that dies holding host memory has it returned by its parent, and a stuck accounting lock fails calls with a retryable 503 instead of hanging. **`/readyz` `host_memory` keys changed**: `large_fits_active`/`large_fits_waiting` are now `calls_active`/`calls_waiting`, plus `reserved_bytes` and `calls_offloading`.
+- **TabICL's attention blocks are compiled** (#36, `TABCTX_COMPILE`, default on; `backends/tabicl_compile.py`). Inductor fuses the memory-bound layer norms, residual adds and GELU MLP around each attention kernel: on an A100 the fit's GPU half is 1.7-2.2x faster (10k x 50: 0.91 -> 0.52 s; 40k x 100: 7.65 -> 4.30 s; 20k x 200: 6.51 -> 2.92 s) and its peak ~20% lower (40k x 100: 33.1 -> 27.0 GiB); 100-row predicts 20-35% faster. Accuracy, AUC and R^2 are unchanged on the evaluation datasets, and predicts beside a fit stay correct. Ray Serve, 6 mixed users on 4x A100: 5.78 -> 6.99 predicts/s, small-fit p95 1.48 -> 0.95 s. tabicl's ssmax log(n) is now computed symbolically inside compiled code, so a new table size doesn't recompile. Each worker's compile pool is sized to its share of the container's CPUs (`TORCHINDUCTOR_COMPILE_THREADS`, when unset). Only tabicl's fp16 path (1,024+ rows) is compiled; its common graphs (~50) compile in `warmup()` at replica start (~130 s per worker, workers in parallel); shapes it doesn't cover (very large tables) compile on first use, ~10 s once per worker. Compile errors fall back to eager, with a warning after warmup if nothing compiled. The lower peak doesn't lower what admission must reserve: the least free memory a compiled fit succeeds with is the same as eager's or up to ~2 GiB more (80k x 10: 22.9 vs 20.7 GiB), so compiled TabICL is admitted against its own requirement grid (per shape, the larger of the two measurements; `minfree-tabicl-kv-compiled-a100.jsonl`).
+- **Host-memory admission** (new; `memory/host.py`). Large TabICL fits offload stage outputs to host RAM (~46 GiB at 10M cells); two at once could get the pod OOM-killed with every tenant's context. Each backend estimates a call's host use (TabICL ~5 KB per cell when it offloads, else nothing; TabPFN ~0.6 KB per cell), and a guard admits calls against one budget shared by every worker process on the node: first come first served (a large fit can't be starved by modest ones), at most `slots` offloading calls at once (each budgets 1/slots of free host memory for its staging), calls under 1 GiB never wait. The budget is the container's memory limit x 0.85 minus its working set (every process in the pod) and an 8 GiB reserve, re-measured once the per-GPU workers have warmed up. A worker that dies holding host memory has it returned by its parent; a stuck accounting lock fails calls with `HostMemoryUnavailableError` (new, exported; HTTP 503, retryable) instead of hanging. `/readyz` gains `host_memory` (`budget_bytes`, `reserved_bytes`, `calls_active`, `calls_offloading`, `calls_waiting`, `large_fit_slots`).
 - **TabICL fits prepare their CPU half before taking the device**: preprocessing (~30-37% of a fit) overlaps another fit's GPU half -- back-to-back fits 17-26% faster on an A100. At most one fit per device prepares ahead, so queued fits don't each hold preprocessed copies in host memory.
 - Worker processes shut down together under one 10 s deadline (was up to 10 s each, past a pod's grace period).
 - **tabicl's per-call log(n) host-to-device syncs are cached** (bit-identical): 10k x 50 fits 1.98 -> 1.66 s on an A100.
diff --git a/README.md b/README.md
index ebd02a6..ef5f683 100644
--- a/README.md
+++ b/README.md
@@ -254,6 +254,11 @@ behind the same API. No client changes and no routing header needed.
 
 ## Large tables: upload, then fit by reference
 
+Tables too big for GPU and host memory together can still fit: TabICL
+offloads stage outputs to memory-mapped files under `TABCTX_OFFLOAD_DIR`
+(default `<TABCTX_SPILL_DIR>/offload`), slower but without exhausting the
+pod's RAM (one A100 fit 130k x 500 this way).
+
 Inline JSON tables are convenient to ~10^4 rows and painful beyond.
 For big tables, stream a CSV first and reference it (the same
 upload/orchestration split hosted tabular APIs use, minus the object
@@ -395,6 +400,10 @@ pip install "tabctx[dev]"        # + test dependencies
   (`TABCTX_KV_CACHE=kv|repr|off`), with predictions verified identical
   to the uncached path -- and loads the pretrained backbone once per
   process instead of per fit.
+- **int8 context cache** (`TABCTX_KV_PRECISION=int8`, the default):
+  TabICL's ICL key/value cache, most of a context, is stored in int8:
+  about half the memory per context, identical labels and metrics on the
+  evaluation sets, probabilities within 0.009 (`native` for exact).
 - **Compiled attention blocks** (`TABCTX_COMPILE=1`, the default):
   TabICL's attention blocks run through `torch.compile`, which fuses the
   memory-bound layer norms and residual adds around each attention
diff --git a/ROADMAP.md b/ROADMAP.md
index 7bf5ef4..fbb4f5b 100644
--- a/ROADMAP.md
+++ b/ROADMAP.md
@@ -130,8 +130,8 @@ no FP8).
     table sizes (9-82 s each, even with its chunked-SDPA loop kept out of
     the graph), so it stays off until those are traced down.
 14. **Serving-framework-agnostic core** ([#39](https://github.com/VectorInstitute/tabctx/issues/39)): the router, per-GPU workers
-    and scheduling are plain Python with no Ray dependency; a Dynamo
-    worker hosts the same core. Placement is a pluggable policy
+    and scheduling are plain Python with no Ray dependency, so a Dynamo
+    worker can host the same core. Placement is a pluggable policy
     (`placement.py`: the router hands it a lock-free view of each device)
     and the router emits residency events (`events.py`: placed/removed,
     to any sink; `/v1/tabctx/residency` serves the latest), so a platform
diff --git a/gke-out/minfree-tabicl-kv-compiled-a100.jsonl b/gke-out/minfree-tabicl-kv-compiled-a100.jsonl
deleted file mode 100644
index 864a2c4..0000000
--- a/gke-out/minfree-tabicl-kv-compiled-a100.jsonl
+++ /dev/null
@@ -1,16 +0,0 @@
-{"rows": 2000, "features": 10, "status": "ok", "resident_gib": 0.39, "unconstrained_peak_gib": 0.59, "unconstrained_fit_s": 0.184, "min_free_gib": 0.69, "min_working_gib": 0.3, "fit_s_at_min": 0.199, "max_slowdown": null}
-{"rows": 20000, "features": 10, "status": "ok", "resident_gib": 3.68, "unconstrained_peak_gib": 5.53, "unconstrained_fit_s": 11.418, "min_free_gib": 5.6, "min_working_gib": 1.92, "fit_s_at_min": 0.956, "max_slowdown": null}
-{"rows": 45000, "features": 10, "status": "ok", "resident_gib": 8.27, "unconstrained_peak_gib": 12.41, "unconstrained_fit_s": 8.189, "min_free_gib": 12.55, "min_working_gib": 4.28, "fit_s_at_min": 3.269, "max_slowdown": null}
-{"rows": 80000, "features": 10, "status": "ok", "resident_gib": 14.68, "unconstrained_peak_gib": 22.04, "unconstrained_fit_s": 13.484, "min_free_gib": 22.94, "min_working_gib": 8.26, "fit_s_at_min": 8.573, "max_slowdown": null}
-{"rows": 2000, "features": 50, "status": "ok", "resident_gib": 0.45, "unconstrained_peak_gib": 1.09, "unconstrained_fit_s": 0.242, "min_free_gib": 1.05, "min_working_gib": 0.6, "fit_s_at_min": 0.265, "max_slowdown": null}
-{"rows": 10000, "features": 50, "status": "ok", "resident_gib": 1.94, "unconstrained_peak_gib": 5.14, "unconstrained_fit_s": 0.688, "min_free_gib": 3.09, "min_working_gib": 1.15, "fit_s_at_min": 0.718, "max_slowdown": null}
-{"rows": 40000, "features": 50, "status": "ok", "resident_gib": 7.4, "unconstrained_peak_gib": 17.6, "unconstrained_fit_s": 14.459, "min_free_gib": 11.8, "min_working_gib": 4.4, "fit_s_at_min": 11.828, "max_slowdown": null}
-{"rows": 2000, "features": 100, "status": "ok", "resident_gib": 0.52, "unconstrained_peak_gib": 1.93, "unconstrained_fit_s": 0.319, "min_free_gib": 1.12, "min_working_gib": 0.6, "fit_s_at_min": 1.017, "max_slowdown": null}
-{"rows": 20000, "features": 100, "status": "ok", "resident_gib": 3.81, "unconstrained_peak_gib": 16.08, "unconstrained_fit_s": 12.789, "min_free_gib": 5.99, "min_working_gib": 2.18, "fit_s_at_min": 8.869, "max_slowdown": null}
-{"rows": 50000, "features": 100, "status": "ok", "resident_gib": 9.33, "unconstrained_peak_gib": 24.02, "unconstrained_fit_s": 20.003, "min_free_gib": 14.38, "min_working_gib": 5.05, "fit_s_at_min": 29.394, "max_slowdown": null}
-{"rows": 1000, "features": 300, "status": "ok", "resident_gib": 0.63, "unconstrained_peak_gib": 2.86, "unconstrained_fit_s": 0.422, "min_free_gib": 1.23, "min_working_gib": 0.6, "fit_s_at_min": 1.691, "max_slowdown": null}
-{"rows": 5000, "features": 300, "status": "ok", "resident_gib": 1.38, "unconstrained_peak_gib": 11.43, "unconstrained_fit_s": 6.469, "min_free_gib": 2.26, "min_working_gib": 0.88, "fit_s_at_min": 7.077, "max_slowdown": null}
-{"rows": 1000, "features": 500, "status": "ok", "resident_gib": 0.92, "unconstrained_peak_gib": 4.68, "unconstrained_fit_s": 0.626, "min_free_gib": 1.51, "min_working_gib": 0.59, "fit_s_at_min": 2.475, "max_slowdown": null}
-{"rows": 5000, "features": 500, "status": "ok", "resident_gib": 1.68, "unconstrained_peak_gib": 16.45, "unconstrained_fit_s": 12.621, "min_free_gib": 2.84, "min_working_gib": 1.16, "fit_s_at_min": 12.255, "max_slowdown": null}
-{"rows": 500, "features": 1000, "status": "ok", "resident_gib": 1.56, "unconstrained_peak_gib": 5.36, "unconstrained_fit_s": 1.085, "min_free_gib": 2.43, "min_working_gib": 0.87, "fit_s_at_min": 2.469, "max_slowdown": null}
-{"rows": 2000, "features": 1000, "status": "ok", "resident_gib": 1.84, "unconstrained_peak_gib": 15.37, "unconstrained_fit_s": 12.801, "min_free_gib": 2.99, "min_working_gib": 1.15, "fit_s_at_min": 8.706, "max_slowdown": null}
diff --git a/src/tabctx/cache/manager.py b/src/tabctx/cache/manager.py
index 176d0c6..02581ff 100644
--- a/src/tabctx/cache/manager.py
+++ b/src/tabctx/cache/manager.py
@@ -8,13 +8,10 @@
 analog of what PagedAttention did for per-request LLM KV caches, scoped here
 to per-training-set tabular ICL contexts.
 
-v1 scaling limit (see tabctx repo README/roadmap): one RLock guards the
-entire cache, and the engine holds it across backend calls too, so all GPU
-work on a replica is effectively serialized. That's the right tradeoff on a
-single GPU under real memory pressure -- two concurrent fits/predicts against
-the same device would just contend for the same memory anyway -- but it caps
-throughput at one in-flight backend call per replica until a future version
-does finer-grained locking or real multi-replica routing.
+Locking. The cache's lock guards only its bookkeeping (entries, usage,
+pins, ids in transit); spill reads and writes run outside it. Scheduling
+GPU work -- one fit at a time per device, light predicts beside it -- is
+the engine's job (engine.py), so the cache never waits on a backend call.
 """
 
 from __future__ import annotations
@@ -337,11 +334,3 @@ def stats(self) -> EngineStats:
                 free_bytes=self.free_bytes,
                 capacity_bytes=self._capacity_bytes,
             )
-
-    @property
-    def lock(self) -> threading.RLock:
-        """Exposed so the engine can hold one lock across the full
-        check-evict-fit-insert sequence (see module docstring) rather than
-        risking two requests both passing admission control against the
-        same free_bytes snapshot and then both calling into the backend."""
-        return self._lock
diff --git a/src/tabctx/memory/host.py b/src/tabctx/memory/host.py
index 1ca3bf0..b120244 100644
--- a/src/tabctx/memory/host.py
+++ b/src/tabctx/memory/host.py
@@ -72,12 +72,25 @@
 # slots for its waiting calls (a worker serves at most 16 calls at once;
 # an in-process guard shared by several devices' engines needs more).
 QUEUE_SLOTS = 32
+# Shared state header: next ticket, budget bytes, large-fit slots.
+_HEADER = 3
 
 
 CGROUP_LIMIT_PATHS = (
     "/sys/fs/cgroup/memory.max",
     "/sys/fs/cgroup/memory/memory.limit_in_bytes",
 )
+# (usage, stat) per cgroup version: the container's working set is usage
+# minus inactive page cache (what the kernel reclaims first; the kubelet's
+# eviction signal is computed the same way).
+CGROUP_USAGE_PATHS = (
+    ("/sys/fs/cgroup/memory.current", "/sys/fs/cgroup/memory.stat", "inactive_file"),
+    (
+        "/sys/fs/cgroup/memory/memory.usage_in_bytes",
+        "/sys/fs/cgroup/memory/memory.stat",
+        "total_inactive_file",
+    ),
+)
 
 
 def detect_host_budget_bytes(
@@ -86,16 +99,41 @@ def detect_host_budget_bytes(
     limit_paths: tuple[str, ...] = CGROUP_LIMIT_PATHS,
     meminfo_path: str = "/proc/meminfo",
     status_path: str = "/proc/self/status",
+    usage_paths: tuple[tuple[str, str, str], ...] = CGROUP_USAGE_PATHS,
 ) -> int | None:
-    """Host memory available to large calls: this process's memory limit
-    (cgroup v2/v1, else physical RAM) x fraction, minus what the process
-    already holds (CUDA contexts, model weights) and a reserve for the
-    rest of serving (uploads, contexts' host state). None when the limit
+    """Host memory available to large calls: the memory limit (cgroup
+    v2/v1, else physical RAM) x fraction, minus what is already in use and
+    a reserve for the rest of serving (uploads, contexts' host state). In
+    use is the container's working set when the cgroup reports it -- every
+    process in the pod: per-GPU workers' models and CUDA contexts, the
+    serving framework -- else this process's RSS. None when the limit
     can't be read."""
     limit = _read_limit(limit_paths, meminfo_path)
     if limit is None:
         return None
-    return max(0, int(limit * fraction) - _rss_bytes(status_path) - reserve_bytes)
+    in_use = _working_set_bytes(usage_paths)
+    if in_use is None:
+        in_use = _rss_bytes(status_path)
+    return max(0, int(limit * fraction) - in_use - reserve_bytes)
+
+
+def _working_set_bytes(usage_paths: tuple[tuple[str, str, str], ...]) -> int | None:
+    for usage_path, stat_path, inactive_key in usage_paths:
+        try:
+            usage = int(Path(usage_path).read_text().strip())
+        except (OSError, ValueError):
+            continue
+        inactive = 0
+        try:
+            for line in Path(stat_path).read_text().splitlines():
+                key, _, value = line.partition(" ")
+                if key == inactive_key:
+                    inactive = int(value)
+                    break
+        except (OSError, ValueError):
+            pass
+        return max(0, usage - inactive)
+    return None
 
 
 def _rss_bytes(status_path: str) -> int:
@@ -129,7 +167,7 @@ class HostMemoryGuard:
     """Node- or process-wide host memory admission (see module docstring).
 
     Shared state (a plain list in-process, an mp Array across worker
-    processes): ``[next_ticket]`` then, per owner, ``[reserved bytes,
+    processes): ``[next_ticket, budget, slots]`` then, per owner, ``[reserved bytes,
     active calls, offloading calls, queue slots...]``. A waiting call
     parks its ticket in one of its owner's queue slots; the oldest parked
     ticket is the head of the queue.
@@ -146,24 +184,14 @@ def __init__(
         lock_timeout_s: float = LOCK_TIMEOUT_S,
         queue_slots: int = QUEUE_SLOTS,
     ) -> None:
-        self.budget_bytes = budget_bytes
         self.large_fit_bytes = large_fit_bytes
-        self.slots = max(1, budget_bytes // large_fit_bytes)
-        if budget_bytes < large_fit_bytes:
-            logger.warning(
-                "host memory budget %.1f GiB is below one large fit's measured "
-                "peak (%.0f GiB): large tables may exhaust this pod; give it more "
-                "memory or rely on disk offload (TABCTX_OFFLOAD_DIR)",
-                budget_bytes / GiB,
-                large_fit_bytes / GiB,
-            )
         self._large_fit_cells = large_fit_cells
         self._min_guarded = min_guarded_bytes
         self._lock_timeout = lock_timeout_s
         self._owners = max(1, owners)
         self._owner = 0
         self._block_size = 3 + max(1, queue_slots)
-        size = 1 + self._owners * self._block_size
+        size = _HEADER + self._owners * self._block_size
         if mp_context is None:
             # In-process: a condition variable (no process can die holding it).
             self._cond: Any = threading.Condition()
@@ -177,6 +205,43 @@ def __init__(
             self._cond = None
             self._lock = mp_context.Lock()
             self._state = mp_context.Array("q", size, lock=False)
+        self._apply_budget(budget_bytes)
+
+    @property
+    def budget_bytes(self) -> int:
+        return int(self._state[1])
+
+    @property
+    def slots(self) -> int:
+        """Offloading calls admitted at once: large fits the budget holds."""
+        return int(self._state[2])
+
+    def set_budget(self, budget_bytes: int) -> None:
+        """Re-budget, e.g. once per-GPU workers have started and warmed up
+        and their own memory is known; shared with every process."""
+        try:
+            with self._locked():
+                self._apply_budget(budget_bytes)
+                self._notify()
+        except HostMemoryUnavailableError:
+            logger.error("host memory budget not updated (lock unavailable)")
+
+    def _apply_budget(self, budget_bytes: int) -> None:
+        self._state[1] = budget_bytes
+        self._state[2] = max(1, budget_bytes // self.large_fit_bytes)
+        logger.info(
+            "host memory budget %.1f GiB (%d large-fit slots)",
+            budget_bytes / GiB,
+            self._state[2],
+        )
+        if budget_bytes < self.large_fit_bytes:
+            logger.warning(
+                "host memory budget %.1f GiB is below one large fit's measured "
+                "peak (%.0f GiB): large tables may exhaust this pod; give it more "
+                "memory or rely on disk offload (TABCTX_OFFLOAD_DIR)",
+                budget_bytes / GiB,
+                self.large_fit_bytes / GiB,
+            )
 
     def bind(self, owner: int) -> None:
         """Account this process's calls to `owner` (a worker's index), so
@@ -305,7 +370,7 @@ def _stuck(self) -> HostMemoryUnavailableError:
         )
 
     def _block(self, owner: int) -> int:
-        return 1 + owner * self._block_size
+        return _HEADER + owner * self._block_size
 
     def _park(self, base: int, ticket: int) -> int | None:
         for i in range(base + 3, base + self._block_size):
diff --git a/src/tabctx/serve/app.py b/src/tabctx/serve/app.py
index 743af29..958a6b7 100644
--- a/src/tabctx/serve/app.py
+++ b/src/tabctx/serve/app.py
@@ -54,6 +54,7 @@
     TenantRequiredError,
     resolve_tenant_id,
     scope_dataset_id,
+    visible_dataset_id,
 )
 from tabctx.serve.uploads import UploadStore
 from tabctx.serve.workers import WorkerDiedError
@@ -278,7 +279,10 @@ def _legacy_predict_sync(self, req: LegacyPredictRequest) -> LegacyPredictRespon
         n_train, n_test = len(req.train_X), len(req.test_X)
         n_features = len(req.train_X[0]) if n_train else 0
 
-        gpus = self._cuda_devices() if cuda else []
+        # Peak GPU memory of this call, measured in-process. With worker
+        # processes the call ran elsewhere: this process holds no models,
+        # and touching each GPU would only create CUDA contexts here.
+        gpus = self._cuda_devices() if cuda and not self._workers else []
         for d in gpus:
             torch.cuda.reset_peak_memory_stats(d)
         start = time.monotonic()
@@ -532,17 +536,35 @@ def models(self):
         }
 
     @fastapi_app.get("/v1/tabctx/residency")
-    def residency(self):
-        """Where contexts are now, per device, and the router's latest
-        placements and removals (multi-GPU replicas): for a platform
-        scheduler or an operator. Events are advisory (contexts dropped by
-        capacity eviction without spill aren't reported); `devices` is the
-        current truth to reconcile against. A single-GPU replica reports
-        no per-device view or events."""
+    def residency(self, request: Request):
+        """Where the caller's contexts are now, per device, and the router's
+        latest placements and removals of them (multi-GPU replicas). Scoped
+        like every /v1/tabctx route: a tenant sees only its own datasets.
+        Events are advisory (contexts dropped by capacity eviction without
+        spill aren't reported); `devices` is the current truth to
+        reconcile against. A single-GPU replica reports neither."""
+        try:
+            tenant_id = resolve_tenant_id(dict(request.headers))
+        except _AUTH_ERRORS as e:
+            raise _map_error(e) from e
+
+        def mine(ids: list[str]) -> list[str]:
+            visible = (visible_dataset_id(tenant_id, i) for i in ids)
+            return [i for i in visible if i is not None]
+
         residency = getattr(self._engine, "residency", None)
+        devices = residency() if callable(residency) else {}
+        events = []
+        for event in self._residency.snapshot():
+            shown = visible_dataset_id(tenant_id, event["dataset_id"])
+            if shown is not None:
+                events.append({**event, "dataset_id": shown})
         return {
-            "devices": residency() if callable(residency) else {},
-            "events": self._residency.snapshot(),
+            "devices": {
+                device: {kind: mine(ids) for kind, ids in view.items()}
+                for device, view in devices.items()
+            },
+            "events": events,
         }
 
     @fastapi_app.get("/v1/tabctx/limits")
diff --git a/src/tabctx/serve/factory.py b/src/tabctx/serve/factory.py
index 4e36dbe..e4949b1 100644
--- a/src/tabctx/serve/factory.py
+++ b/src/tabctx/serve/factory.py
@@ -766,6 +766,12 @@ def _build_process_engine(
         host_guard=host_guard,
         mp_context=ctx,
     )
+    if host_guard is not None:
+        # Measured before the workers existed; now that each holds its
+        # models, CUDA context and compiled kernels, budget what is left.
+        rebudget = detect_host_budget_bytes()
+        if rebudget is not None:
+            host_guard.set_budget(rebudget)
     # This process holds no backends: /limits and /readyz use estimators
     # built from the same settings (calibration, not this replica's
     # learned observations).
diff --git a/src/tabctx/serve/tenancy.py b/src/tabctx/serve/tenancy.py
index ef96835..7e50d79 100644
--- a/src/tabctx/serve/tenancy.py
+++ b/src/tabctx/serve/tenancy.py
@@ -146,3 +146,13 @@ def unscope_dataset_id(tenant_id: str | None, scoped_id: str) -> str:
             f"scoped id {scoped_id!r} does not belong to tenant {tenant_id!r}"
         )
     return scoped_id[len(prefix) :]
+
+
+def visible_dataset_id(tenant_id: str | None, scoped_id: str) -> str | None:
+    """The caller-facing id of a cache id the caller may see, else None: a
+    tenant sees only its own datasets; a caller without a tenant sees only
+    unscoped ones (never another tenant's, whose ids embed the tenant)."""
+    if tenant_id is None:
+        return None if scoped_id.startswith(f"{_SCOPE_PREFIX}:") else scoped_id
+    prefix = f"{_SCOPE_PREFIX}:{tenant_id}:"
+    return scoped_id[len(prefix) :] if scoped_id.startswith(prefix) else None
diff --git a/src/tabctx/serve/workers.py b/src/tabctx/serve/workers.py
index 219f947..510cb49 100644
--- a/src/tabctx/serve/workers.py
+++ b/src/tabctx/serve/workers.py
@@ -41,6 +41,7 @@
 import itertools
 import logging
 import multiprocessing.util  # noqa: F401 -- registers its exit handler first
+import os
 import pickle
 import threading
 import time
@@ -48,6 +49,7 @@
 from concurrent.futures import Future, ThreadPoolExecutor
 from concurrent.futures import TimeoutError as FutureTimeoutError
 from dataclasses import dataclass
+from pathlib import Path
 from typing import Any
 
 from tabctx.errors import BackendComputeError
@@ -169,6 +171,31 @@ def _report(shard: Any, device: str) -> dict:
     }
 
 
+def available_cpus() -> int:
+    """CPUs this container may use: its cgroup CPU quota (v2 cpu.max) if
+    set, else the CPUs it is scheduled on."""
+    try:
+        quota, period = Path("/sys/fs/cgroup/cpu.max").read_text().split()
+        if quota != "max":
+            return max(1, int(int(quota) / int(period)))
+    except (OSError, ValueError):
+        pass
+    try:
+        return len(os.sched_getaffinity(0))
+    except AttributeError:  # not Linux
+        return os.cpu_count() or 1
+
+
+def _share_cpus(n_workers: int) -> None:
+    """Size this worker's torch.compile pool to its share of the pod's CPUs
+    (Inductor sizes it from the host's cores, not the container's): four
+    workers compiling at once at startup would otherwise oversubscribe the
+    pod and stretch warmup. A deployment's own setting wins."""
+    if "TORCHINDUCTOR_COMPILE_THREADS" not in os.environ:
+        share = max(1, available_cpus() // max(1, n_workers))
+        os.environ["TORCHINDUCTOR_COMPILE_THREADS"] = str(share)
+
+
 def _worker_main(
     conn: Any,
     settings: Any,
@@ -185,6 +212,7 @@ def _worker_main(
 
     if host_guard is not None:
         host_guard.bind(index)  # so the parent can release it if we die
+    _share_cpus(n_devices)
     try:
         shard = build_shard(
             settings,
diff --git a/tests/unit/test_host_memory.py b/tests/unit/test_host_memory.py
index d8a6436..824d109 100644
--- a/tests/unit/test_host_memory.py
+++ b/tests/unit/test_host_memory.py
@@ -387,10 +387,49 @@ def test_budget_from_cgroup_v2_minus_baseline_rss(tmp_path):
         limit_paths=(str(tmp_path / "memory.max"),),
         meminfo_path=str(tmp_path / "none"),
         status_path=str(tmp_path / "status"),
+        usage_paths=(),  # no cgroup usage: this process's RSS
     )
     assert budget == int(64 * GiB * 0.85) - 4 * GiB - 8 * GiB
 
 
+def test_budget_subtracts_the_containers_working_set(tmp_path):
+    """Every process in the pod counts (per-GPU workers, the serving
+    framework), not just this one; inactive page cache doesn't."""
+    from tabctx.memory.host import detect_host_budget_bytes
+
+    (tmp_path / "memory.max").write_text(f"{64 * GiB}\n")
+    (tmp_path / "memory.current").write_text(f"{20 * GiB}\n")
+    (tmp_path / "memory.stat").write_text(f"anon 1\ninactive_file {4 * GiB}\n")
+    (tmp_path / "status").write_text("VmRSS:\t 1048576 kB\n")  # 1 GiB: ignored
+    budget = detect_host_budget_bytes(
+        fraction=0.85,
+        reserve_bytes=8 * GiB,
+        limit_paths=(str(tmp_path / "memory.max"),),
+        meminfo_path=str(tmp_path / "none"),
+        status_path=str(tmp_path / "status"),
+        usage_paths=(
+            (
+                str(tmp_path / "memory.current"),
+                str(tmp_path / "memory.stat"),
+                "inactive_file",
+            ),
+        ),
+    )
+    assert budget == int(64 * GiB * 0.85) - 16 * GiB - 8 * GiB
+
+
+def test_rebudgeting_is_shared_with_worker_processes():
+    import multiprocessing as mp
+
+    guard = HostMemoryGuard(
+        100 * GiB, large_fit_bytes=48 * GiB, mp_context=mp.get_context("spawn")
+    )
+    assert (guard.budget_bytes, guard.slots) == (100 * GiB, 2)
+    guard.set_budget(60 * GiB)  # workers' own memory now counted
+    assert (guard.budget_bytes, guard.slots) == (60 * GiB, 1)
+    assert guard.stats()["budget_bytes"] == 60 * GiB
+
+
 def test_unlimited_cgroup_falls_back_to_physical_memory(tmp_path):
     from tabctx.memory.host import detect_host_budget_bytes
 
@@ -403,6 +442,7 @@ def test_unlimited_cgroup_falls_back_to_physical_memory(tmp_path):
         limit_paths=(str(tmp_path / "v2"), str(tmp_path / "v1")),
         meminfo_path=str(tmp_path / "meminfo"),
         status_path=str(tmp_path / "none"),
+        usage_paths=(),
     )
     assert budget == 16 * GiB
 
diff --git a/tests/unit/test_serve_tenancy.py b/tests/unit/test_serve_tenancy.py
index eaece00..4a84697 100644
--- a/tests/unit/test_serve_tenancy.py
+++ b/tests/unit/test_serve_tenancy.py
@@ -89,3 +89,15 @@ def test_unscope_foreign_id_refused(self):
         scoped = scope_dataset_id("acme", "ds-1")
         with pytest.raises(ValueError):
             unscope_dataset_id("other", scoped)
+
+
+def test_callers_see_only_their_own_dataset_ids():
+    from tabctx.serve.tenancy import scope_dataset_id, visible_dataset_id
+
+    acme = scope_dataset_id("acme", "churn")
+    other = scope_dataset_id("globex", "secret")
+    assert visible_dataset_id("acme", acme) == "churn"
+    assert visible_dataset_id("acme", other) is None
+    assert visible_dataset_id("acme", "unscoped") is None
+    assert visible_dataset_id(None, "unscoped") == "unscoped"
+    assert visible_dataset_id(None, acme) is None  # never another tenant's
diff --git a/tests/unit/test_workers.py b/tests/unit/test_workers.py
index 68fbb9d..261965e 100644
--- a/tests/unit/test_workers.py
+++ b/tests/unit/test_workers.py
@@ -208,3 +208,17 @@ def test_mp_exit_handler_is_registered_before_ours():
         check=True,
     )
     assert out.stdout.strip() == "True"
+
+
+def test_workers_share_the_pods_cpus_for_compiling(monkeypatch):
+    import tabctx.serve.workers as workers
+
+    monkeypatch.delenv("TORCHINDUCTOR_COMPILE_THREADS", raising=False)
+    monkeypatch.setattr(workers, "available_cpus", lambda: 30)
+    workers._share_cpus(4)
+    import os
+
+    assert os.environ["TORCHINDUCTOR_COMPILE_THREADS"] == "7"
+    monkeypatch.setenv("TORCHINDUCTOR_COMPILE_THREADS", "2")  # deployer's wins
+    workers._share_cpus(4)
+    assert os.environ["TORCHINDUCTOR_COMPILE_THREADS"] == "2"

From b13c9dd1f180c399fbb73b7e77b202164ffe3719 Mon Sep 17 00:00:00 2001
From: Amrit Krishnan <amrit110@gmail.com>
Date: Sun, 18 Oct 2026 02:35:48 -0400
Subject: [PATCH 38/39] fix(engine): make room against the device's real free
 memory too

On the 4x L4 rehearsal an admitted 250 x 1000 TabPFN fit OOMed on a
22 GB card holding 18 GB: memory outside the cache's books grows after
warmup (allocator fragmentation, compiled kernels' workspaces), so the
book headroom overstated what was free. _make_room now also reads the
device's real available memory (free + torch's reusable cache;
_cuda.device_available_bytes, wired per device by the factory) and evicts
until both agree. A fit retrying after an OOM frees contexts parked during
its failed attempt (explicit evicts while it measured) before retrying.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 CHANGELOG.md                         |  1 +
 src/tabctx/backends/_cuda.py         | 14 ++++++
 src/tabctx/engine.py                 | 28 ++++++++++--
 src/tabctx/serve/factory.py          |  5 +++
 tests/unit/test_engine_scheduling.py | 66 ++++++++++++++++++++++++++++
 5 files changed, 110 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index a33da89..0347d54 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -30,6 +30,7 @@
 - **Weights load at replica start** (`warmup()` on both backends, called by the factory) instead of inside the first tenant's fit, which reported the checkpoint as part of its context (a 200x6 TabPFN context was booked at 855 MB, TabICL's at 226 MB) and paid the load latency.
 
 ### Fixed
+- **Fits OOMed on GPUs the books said had room** (4x L4 rehearsal: a 250 x 1000 TabPFN fit beside 18 GB of contexts on a 22 GB card). Memory outside the cache's accounting grows after warmup (allocator fragmentation, compiled kernels' workspaces), so making room now also checks the device's real free memory (free plus torch's reusable cache) and evicts until both agree. A fit retrying after an OOM first frees contexts that were parked during its failed attempt.
 - **Router 404 for a context being spilled** (4xL4 rehearsal: a fit returned 200 and the next predict 404'd 27 ms later): a context moving between the GPU cache and the spill tier was listed by neither; it is now tracked in transit, and predict re-locates once after a miss.
 - **TabICL shared-backbone cache was per-module, not per-call**: interleaved calls could read another context's cache. It is now thread-local.
 - A device OOM while restoring a spilled context no longer deletes the spilled copy.
diff --git a/src/tabctx/backends/_cuda.py b/src/tabctx/backends/_cuda.py
index 67a4a12..c6f73ca 100644
--- a/src/tabctx/backends/_cuda.py
+++ b/src/tabctx/backends/_cuda.py
@@ -40,6 +40,20 @@ def release_cached_device_memory() -> None:
         torch.cuda.empty_cache()
 
 
+def device_available_bytes(device: str) -> int | None:
+    """Device memory a new allocation can use now: free on the device plus
+    what torch's caching allocator holds reserved but unallocated (reused
+    before asking the driver). None off CUDA."""
+    if not is_cuda(device):
+        return None
+    import torch
+
+    d = torch.device(device)
+    free, _ = torch.cuda.mem_get_info(d)
+    cached = torch.cuda.memory_reserved(d) - torch.cuda.memory_allocated(d)
+    return int(free + max(0, cached))
+
+
 def device_scope(device: str) -> contextlib.AbstractContextManager[Any]:
     """Make ``device`` torch's current CUDA device for the block; a no-op
     on CPU."""
diff --git a/src/tabctx/engine.py b/src/tabctx/engine.py
index 6c9120f..566013a 100644
--- a/src/tabctx/engine.py
+++ b/src/tabctx/engine.py
@@ -10,7 +10,7 @@
 import gc
 import threading
 import uuid
-from collections.abc import Iterator
+from collections.abc import Callable, Iterator
 
 from tabctx.backends._cuda import release_cached_device_memory
 from tabctx.backends.base import TabularICLBackend
@@ -156,6 +156,7 @@ def __init__(
         default_backend: str | None = None,
         host_guard: HostMemoryGuard | None = None,
         device_baseline_bytes: int = 0,
+        device_available: Callable[[], int | None] | None = None,
     ) -> None:
         """Single-backend form: TabctxEngine(backend=, cache=, estimator=).
         Multi-backend form: TabctxEngine(backends={name: b}, cache=,
@@ -211,6 +212,12 @@ def __init__(
         # 22 GB L4 the 15% transient reserve doesn't absorb it: fits the
         # books said were safe OOMed).
         self._baseline_bytes = device_baseline_bytes
+        # The device's real free memory (see _cuda.device_available_bytes):
+        # making room also checks it, since memory outside the books grows
+        # after warmup (allocator fragmentation, compiled kernels'
+        # workspaces) -- on a 22 GB L4 an admitted fit OOMed beside 18 GB
+        # the books didn't fully see.
+        self._device_available = device_available
 
     @property
     def backend_names(self) -> list[str]:
@@ -400,7 +407,10 @@ def _make_room(self, estimator: MemoryEstimator, needed: int) -> None:
         the device) and those contexts go too."""
 
         def short() -> bool:
-            return needed > self._headroom(estimator, self._cache.stats().used_bytes)
+            if needed > self._headroom(estimator, self._cache.stats().used_bytes):
+                return True
+            available = self._device_available() if self._device_available else None
+            return available is not None and needed > available
 
         while short():
             if self._cache.evict_one() is None:
@@ -425,6 +435,11 @@ def _measuring(self) -> Iterator[None]:
                 parked, self._parked = self._parked, []
             del parked
 
+    def _drop_parked(self) -> None:
+        with self._release_lock:
+            parked, self._parked = self._parked, []
+        del parked
+
     def _release(self, context: CachedContext | None) -> None:
         """Drop the caller's last reference to a removed context now, or
         park it until the running fit has measured itself. The caller must
@@ -459,10 +474,15 @@ def attempt() -> object:
             return attempt()
         except BackendComputeError:
             pass
-        gc.collect()
-        release_cached_device_memory()
         with self._predict_lock:
             self._make_room(estimator, estimated)
+            # Contexts removed during the failed attempt (explicit evicts,
+            # re-fits) were parked so as not to skew its measurement
+            # (_measuring); that measurement is void and the retry measures
+            # afresh, so free them now and give the retry their memory.
+            self._drop_parked()
+            gc.collect()
+            release_cached_device_memory()
             return attempt()
 
     def _learn(
diff --git a/src/tabctx/serve/factory.py b/src/tabctx/serve/factory.py
index e4949b1..907ed4f 100644
--- a/src/tabctx/serve/factory.py
+++ b/src/tabctx/serve/factory.py
@@ -69,6 +69,7 @@
 
 from __future__ import annotations
 
+import functools
 import importlib
 import logging
 import os
@@ -77,6 +78,7 @@
 from pathlib import Path
 from typing import Literal
 
+from tabctx.backends._cuda import device_available_bytes
 from tabctx.backends.base import TabularICLBackend
 from tabctx.backends.kv_precision import KV_PRECISIONS, KVPrecision
 from tabctx.cache.manager import ContextCacheManager
@@ -675,6 +677,9 @@ def build_shard(
         default_backend=default,
         host_guard=host_guard,
         device_baseline_bytes=0 if fake_only else measure_device_baseline(device),
+        device_available=None
+        if fake_only
+        else functools.partial(device_available_bytes, device),
     )
     return Shard(
         engine=engine,
diff --git a/tests/unit/test_engine_scheduling.py b/tests/unit/test_engine_scheduling.py
index 779f714..277dfa8 100644
--- a/tests/unit/test_engine_scheduling.py
+++ b/tests/unit/test_engine_scheduling.py
@@ -393,3 +393,69 @@ def test_small_fits_skip_the_prepare_gate(monkeypatch):
         assert engine._cache.holds("small")
     finally:
         engine._prepare_gate.release()
+
+
+class _Payload:
+    def __init__(self, inner):
+        self.inner = inner
+
+
+class EvictDuringFitFake(GatedFake):
+    """Its first fit fails after another request evicted a context
+    mid-fit; records whether that context was freed before the retry."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.engine = None
+        self.watch = None
+        self.freed_before_retry = None
+
+    def fit(self, X, y, task):
+        if self.fail_next_fits:
+            self.fail_next_fits -= 1
+            self.engine.evict("old")  # parked: this fit is measuring
+            raise BackendComputeError("CUDA OOM during fit()")
+        if self.watch is not None:
+            import gc
+
+            gc.collect()
+            self.freed_before_retry = self.watch() is None
+        return _Payload(super().fit(X, y, task))
+
+    def predict(self, payload, X_test, return_proba=False):
+        return super().predict(payload.inner, X_test, return_proba)
+
+
+def test_a_retry_after_oom_frees_contexts_parked_during_the_failed_attempt():
+    """Contexts removed while a fit measures are parked, not freed; the
+    failed attempt's measurement is void, so the retry gets their memory."""
+    import weakref
+
+    backend = EvictDuringFitFake()
+    engine = TabctxEngine(
+        backend=backend, estimator=CellEstimator(), cache=ContextCacheManager(CAPACITY)
+    )
+    backend.engine = engine
+    engine.fit(*table(10), dataset_id="old")
+    backend.watch = weakref.ref(engine._cache.peek("old").payload)
+    backend.fail_next_fits = 1
+    engine.fit(*table(10), dataset_id="new")
+    assert backend.freed_before_retry is True
+
+
+def test_making_room_also_respects_the_devices_real_free_memory():
+    """The books say there's room; the device doesn't: evict anyway."""
+    available = [10**12]
+    engine = TabctxEngine(
+        backend=GatedFake(),
+        estimator=CellEstimator(),
+        cache=ContextCacheManager(CAPACITY),
+        device_available=lambda: available[0],
+    )
+    engine.fit(*table(10), dataset_id="a")
+    engine.fit(*table(10), dataset_id="b")
+    available[0] = 50  # less than a 10-row fit needs (100), whatever the books say
+    engine._device_available = lambda: 10**12 if not engine._cache.holds("a") else 50
+    engine.fit(*table(10), dataset_id="c")
+    assert not engine._cache.holds("a")  # the coldest went
+    assert engine._cache.holds("b") and engine._cache.holds("c")

From b417dd8e6de6e2e2a04b8e26aa1f2a98b632eb5f Mon Sep 17 00:00:00 2001
From: Amrit Krishnan <amrit110@gmail.com>
Date: Sun, 18 Oct 2026 10:25:11 -0400
Subject: [PATCH 39/39] test: import the shared multidevice helpers the way
 CI's pytest path allows; ruff format

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 tests/unit/test_placement_events.py | 2 +-
 tests/unit/test_tabicl_kv.py        | 5 ++++-
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/tests/unit/test_placement_events.py b/tests/unit/test_placement_events.py
index 4a5eaad..3b58fed 100644
--- a/tests/unit/test_placement_events.py
+++ b/tests/unit/test_placement_events.py
@@ -1,11 +1,11 @@
 """Pluggable placement (placement.py) and residency events (events.py)."""
 
 import pytest
+from test_multidevice import build, table, where
 
 from tabctx.events import EventBus, RecentEvents, ResidencyEvent
 from tabctx.multidevice import MultiDeviceEngine
 from tabctx.placement import DeviceView, FitRequest, LeastActiveThatFits
-from tests.unit.test_multidevice import build, table, where
 
 REQ = FitRequest(n_train=10, n_features=1, backend="m")
 
diff --git a/tests/unit/test_tabicl_kv.py b/tests/unit/test_tabicl_kv.py
index fd9dcc1..08f801f 100644
--- a/tests/unit/test_tabicl_kv.py
+++ b/tests/unit/test_tabicl_kv.py
@@ -7,7 +7,10 @@
 
 torch = pytest.importorskip("torch")
 
-from tabctx.backends.tabicl_kv import CompressedKVEntry, compress_icl_cache  # noqa: E402
+from tabctx.backends.tabicl_kv import (  # noqa: E402
+    CompressedKVEntry,
+    compress_icl_cache,
+)
 
 
 @dataclass