[amrit110]

Part of #47 (the predict stall; wide-fit compiles in-request are next).

Problem

On prod (4× A40), predicts against wide contexts took a flat ~24 s (tabicl-v2) and ~11 s (tabpfn-3.5) from 3,500 × 500 up, whatever the table size. Reproduced on an A100 (14.4 s / 6.3 s). cProfile of one predict (benchmarks/profile_predict.py) showed backend.predict running 100 times for one 100-row request: choose_chunk_rows sized chunks with estimate_bytes(n_train, chunk, n_features), which for kv contexts outside the old predict grid fell back to the fit formula. That grows with the training table, so past a size no chunk "fit" and chunking fell to one row.

Fix

• memory/predict.py, PredictCostModel. Measured (benchmarks/calibrate_predict.py, A100, serving config: compiled TabICL, int8 kv), a kv-context predict's peak follows its test rows and their width, plus a small cost per training row (int8 KV dequantized per layer). The model is built from measured marginal costs, not a fitted surface (fitted surfaces underestimated held-out widths by up to a third):

  • per-width cost per test row
  • per-width fixed cost per call
  • cost per training row

  Both width curves are non-decreasing; past the widest measurement they scale in proportion to width. Guaranteed: the bound is at or above every measurement it was built from. Assumed: at most linear growth beyond them (steepest slope between consecutive sizes; affine in width). Over the 68 measurements: TabICL median ×1.16 overestimate, TabPFN ×1.09; held-out widths ≥0.98× before the 1.1 margin.

• AdaptiveMemoryEstimator(predict_cost=…) uses it for kv predicts. repr and off modes, where predicts re-encode the training table, keep the dominance grid. The factory installs it per model from the generated memory/calibration_predict_a100.py.

• choose_chunk_rows, when a call's fixed cost alone exceeds the headroom, sizes only the rows above it against the budget instead of splitting to single rows, which only multiplied that fixed cost. With no headroom (a full device, say beside a fit at its peak) it still uses the smallest chunks, so it adds nothing. MAX_CHUNK_ROWS = LIGHT_PREDICT_ROWS = 2,048, so a light predict is one call.

• Configuration-keyed: the A100 calibration applies only to the configuration it was measured under (kv, int8, compiled TabICL; TabPFN). Other configurations get no model.

Verified

• Unit: 497 pass. The regression test fails 6/6 with the model disabled (checked). New in tests/unit/test_predict_cost.py:
  • the bound holds on synthetic ground truth for unmeasured widths and context sizes;
  • monotone in every dimension;
  • the shipped calibration bounds every measurement;
  • factory wiring for kv vs repr;
  • ragged calibrations: one test size at a width, negative slopes, narrower than measured, convex curves;
  • chunk sizing with a fixed cost over budget, and with no headroom;
  • regression: a 100-row predict against 3,500/10,000/100,000 × 500 fits in the headroom and in one chunk for both models, while the fit model's estimate exceeds it.
• GPU (A100 devbox, benchmarks/bench_stalls.py, engine built as serving builds it, 3 predicts per shape):

Shape	TabICL before → after	TabPFN before → after
3,500 × 500	14.4 s → 0.11 s	6.3 s → 0.15 s
10,000 × 500	14.4 s → 0.11 s	6.3 s → 0.16 s
10,000 × 1,000	— → 0.17–0.21 s	— → 0.23 s
10–100 features	0.05–0.07 s, unchanged	0.06 s, unchanged

Every predict is ≤0.24 s and compiles nothing.

Review

An independent review agent found that the first version's fallback could add an ~18 GB chunk beside a running fit, and that the regression test couldn't fail. Both are fixed in the second commit, along with its should-fixes: configuration keying, docstring accuracy, consecutive slopes, ragged-calibration tests, and the generator rejecting mixed inputs. GPU results were re-run after the fixes, and are unchanged.

Not in this PR

• Wide TabICL fits compile inside the request: 3,500 × 500 builds 3 graphs (12.5 s fit). Warmup coverage is next, in Fixed 10–30 s stalls on fit and predict across table sizes (recompiles?) #47.
• 10,000 × 1,000 fit takes 127 s against 3.5 s at 3,500 × 1,000: the offload cliff, Avoid the offload cliff on large fits (3× slower past ~10M cells) #49.
• Calibration is A100: an A40 calibration is Per-device TabICL calibration grids (L4, A40) for admission #37. TabPFN batches test rows by free memory, so its peak may differ on the A40's 48 GB; the margin and halve-on-OOM cover it until measured.

Also adds benchmarks/sweep_gateway.py, and fixes devbox.sh down for a cluster whose creation failed.

🤖 Generated with Claude Code

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 79.51%. Comparing base (bfac40a) to head (2533f98).

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #54      +/-   ##
==========================================
+ Coverage   78.72%   79.51%   +0.78%     
==========================================
  Files          39       41       +2     
  Lines        3098     3193      +95     
==========================================
+ Hits         2439     2539     +100     
+ Misses        659      654       -5     

Flag	Coverage Δ
integration	30.12% <34.00%> (+0.07%)	⬆️
unit	76.01% <100.00%> (+0.89%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
src/tabctx/chunking.py	100.00% <100.00%> (ø)
src/tabctx/engine.py	97.16% <100.00%> (+0.31%)	⬆️
src/tabctx/memory/adaptive.py	98.86% <100.00%> (+0.05%)	⬆️
src/tabctx/memory/calibration_predict_a100.py	100.00% <100.00%> (ø)
src/tabctx/memory/predict.py	100.00% <100.00%> (ø)
src/tabctx/serve/factory.py	92.92% <100.00%> (+0.30%)	⬆️

... and 1 file with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.