feat: multi-GPU replicas with workload-aware placement (0.12.0)

CHANGELOG.md

@@ -2,6 +2,29 @@
 
 ## [Unreleased]
 
+## [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.
+- `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
+- **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
+- **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.
+- `/readyz` and the legacy endpoint read GPU memory from torch's current device; they now read the replica's own device(s).
+- The TabICL and TabPFN backends measured memory on, and let library internals allocate on, torch's *current* CUDA device rather than their own. Harmless on one GPU; on `cuda:i` it would have charged every context to `cuda:0`. They now measure on their device and run inside its scope.
+
+### Validated
+- Unit (multi-device placement, routing, spill restore, failures) and integration tests (the serve app on two fake devices through Ray). Concurrency tests force interleavings with gates and barriers rather than timing: a fit on an idle device completes while another device is mid-call, lookups and stats never wait on a busy device, concurrent same-id fits leave one copy, a predict during a re-fit is answered. Seven of them fail against the first version of this engine (found by review).
+- On a local fake 3-device replica with 5 concurrent users: with spill on, 0 failed predicts at ~100 predicts/s. Without spill, capacity eviction cost 12 predicts on three devices versus 115 on one.
+
 ## [0.11.0] - 2026-10-17
 ### Changed
 - **The TabPFN model id follows the checkpoint actually loaded.** It was hardcoded to `tabpfn-3`, but `tabpfn` 9.x loads TabPFN-3.5 by default (a different model under its own license, `tabpfn-3-5-license-v1.0`), so a default install served TabPFN-3.5 named `tabpfn-3`. The id is now derived from `tabpfn`'s resolved `model_version` (`TABPFN_MODEL_VERSION`): `v3.5` -> `tabpfn-3.5` (the 9.x default), `v3` -> `tabpfn-3`, `v3.5-fast` -> `tabpfn-3.5-fast`. **Breaking for callers that send `"model": "tabpfn-3"` to a deployment on the default version**: pin `TABPFN_MODEL_VERSION=v3` to keep serving TabPFN-3 under its old id, or switch callers to `tabpfn-3.5`.

README.md

@@ -200,6 +200,42 @@ Details worth knowing:
   torch -- `tests/integration/test_multi_replica_affinity.py` uses it to
   prove the multi-replica contract on a laptop.
 
+## Several GPUs in one replica
+
+A replica that can see several GPUs serves from all of them
+(`TABCTX_DEVICES=auto`, the default, or an explicit `cuda:0,cuda:1,...`):
+one engine per GPU, with its own backends, cache and spill directory,
+behind the same API. No client changes and no routing header needed.
+
+- **Placement.** Each fit goes to a GPU where the table fits *without
+  evicting* anyone, preferring the least active GPU (calls in flight,
+  then calls in the last 30s). If no GPU can take it without eviction,
+  it goes to the GPU with the most headroom. Concurrent users spread
+  across GPUs, and a large table gets a whole GPU's headroom.
+- **Predicts** run on the GPU that holds the context, including after a
+  spill and restore, or a replica restart.
+- **Limits.** Neither TabICL nor TabPFN can split one context across
+  GPUs, so a table's ceiling is one GPU's memory. What more GPUs buy is
+  N x the resident cache and N fits/predicts in flight at once.
+- **Concurrency.** Routing reads lock-free snapshots, so it never waits
+  on a GPU's work: fits and predicts on different GPUs run at the same
+  time. Fits and evicts of the *same* dataset_id are serialized, so a
+  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.
+- **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),
+  `RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=1`.
+- **Spill.** Each GPU spills to its own subdirectory (`cuda_0/`, ...);
+  `TABCTX_SPILL_CAPACITY_BYTES` is the replica's total, split evenly.
+  Changing between one and several GPUs changes the layout, and the
+  replica logs a warning about spilled files it won't restore.
+- `/readyz` reports per-device contexts, bytes, activity and real GPU
+  memory (`real_gpu_memory` keeps its flat fields, summed, plus
+  `per_device`); `/v1/tabctx/limits` reports the best single GPU's
+  headroom and each device's.
+
 ## Large tables: upload, then fit by reference
 
 Inline JSON tables are convenient to ~10^4 rows and painful beyond.

ROADMAP.md

@@ -45,6 +45,44 @@ gets deleted when it empties.
 4. **Docs site** ([#9](https://github.com/VectorInstitute/tabctx/issues/9)): README carries the full story today; a docs site
    matters as external adoption starts.
 
+## Throughput and table size: multi-GPU and CUDA-level work
+
+The LLM serving engines (vLLM, SGLang, TensorRT-LLM) got their gains first
+from memory management and scheduling (paged KV, continuous batching,
+prefix reuse) and then from CUDA kernels. tabctx already has the
+scheduling layer (context cache, usage-aware admission, evict-ahead,
+spill, same-context coalescing). This track adds multi-GPU replicas and
+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,
+   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)).
+
+PyTorch-level first: tiling and chunking, `torch.compile` (generated Triton
+kernels, CUDA graphs via `reduce-overhead`), FlexAttention for paged and
+cross-context attention, SDPA, and `torch.distributed` context
+parallelism. Hand-written Triton/CUDA only where a profile (#31) shows a
+gap those can't close -- the models change too often (TabPFN v3 -> v3.5,
+TabICL v1 -> v2 within months) for version-locked kernels to be the
+default.
+
 ## Standing process note
 
 Every significant fix in this project came from testing against real

benchmarks/bench_copy.py

@@ -0,0 +1,98 @@
+#!/usr/bin/env python3
+"""Context copy bandwidth on a GPU node (#39): what moving a context costs.
+
+Fit/predict disaggregation and rebalancing contexts between GPUs both pay
+for a copy. This measures, for buffer sizes spanning real contexts (MBs to
+tens of GB): GPU->GPU (peer, when available), GPU->pinned host,
+pinned host->GPU, and GPU->pageable host (what pickling a context to the
+spill tier starts with).
+
+Usage: bench_copy.py --sizes-gib 0.25 1 4 16 --out copy-a100x4.json
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+import time
+
+import torch
+
+GiB = 1024**3
+
+
+def bandwidth(fn, nbytes: int, devices: list[torch.device], repeats: int) -> float:
+    fn()  # warm
+    for d in devices:
+        torch.cuda.synchronize(d)
+    t0 = time.perf_counter()
+    for _ in range(repeats):
+        fn()
+    for d in devices:
+        torch.cuda.synchronize(d)
+    return nbytes * repeats / (time.perf_counter() - t0) / GiB
+
+
+def main() -> int:
+    p = argparse.ArgumentParser()
+    p.add_argument("--sizes-gib", type=float, nargs="+", default=[0.25, 1, 4, 16])
+    p.add_argument("--repeats", type=int, default=3)
+    p.add_argument("--out", default=None)
+    args = p.parse_args()
+
+    n = torch.cuda.device_count()
+    gpus = [torch.device(f"cuda:{i}") for i in range(n)]
+    result = {
+        "gpus": [torch.cuda.get_device_name(g) for g in gpus],
+        "peer_access": {
+            f"{i}->{j}": torch.cuda.can_device_access_peer(i, j)
+            for i in range(n)
+            for j in range(n)
+            if i != j
+        },
+        "gib_per_s": [],
+    }
+    for size in args.sizes_gib:
+        nbytes = int(size * GiB)
+        src = torch.empty(nbytes, dtype=torch.uint8, device=gpus[0])
+        pinned = torch.empty(nbytes, dtype=torch.uint8, pin_memory=True)
+        row = {"size_gib": size}
+        row["gpu_to_pinned_host"] = bandwidth(
+            lambda s=src, h=pinned: h.copy_(s, non_blocking=True),
+            nbytes,
+            gpus[:1],
+            args.repeats,
+        )
+        row["pinned_host_to_gpu"] = bandwidth(
+            lambda s=src, h=pinned: s.copy_(h, non_blocking=True),
+            nbytes,
+            gpus[:1],
+            args.repeats,
+        )
+        row["gpu_to_pageable_host"] = bandwidth(
+            lambda s=src: s.cpu(), nbytes, gpus[:1], max(1, args.repeats // 2)
+        )
+        if n > 1:
+            dst = torch.empty(nbytes, dtype=torch.uint8, device=gpus[1])
+            row["gpu0_to_gpu1"] = bandwidth(
+                lambda s=src, d=dst: d.copy_(s, non_blocking=True),
+                nbytes,
+                gpus[:2],
+                args.repeats,
+            )
+            dst = None
+        row = {k: round(v, 2) if isinstance(v, float) else v for k, v in row.items()}
+        result["gib_per_s"].append(row)
+        print(json.dumps(row), flush=True)
+        src = pinned = None
+        torch.cuda.empty_cache()
+    if args.out:
+        with open(args.out, "w") as f:
+            json.dump(result, f, indent=2)
+    print(json.dumps({k: v for k, v in result.items() if k != "gib_per_s"}))
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())