Replay CUDA graphs where they are measured to be worth it

[JingangQu]

Follow-up to #994. Small KumoTabular forwards leave the GPU waiting for kernel launches. This PR replays CUDA graphs for repeated input layouts, e.g., across folds of a table, with the graph cache from #971, but decides per layout from measurements instead of a fixed input size limit.

Changes

Replay CUDA graphs of repeated forwards

• _KumoTabular.forward runs inference through GraphCache on CUDA, without gradients or the key/value cache and outside of torch.compile. Graphs are keyed by input layout, autocast state and parameter storage, and replays are bit-identical to eager execution.

Decide per layout from measurements

• GraphCache times eager calls with CUDA events. Once two of them have finished, it captures a graph into its own torch.cuda.MemPool. The first replay, which also uploads the graph, is not timed.
• A graph is kept if two replays each save a larger share of the fastest eager call's time than the share of device memory taken by its pool. Otherwise it is dropped together with its pool, so graphs of large, compute-bound inputs do not hold on to memory.
• Once a graph is dropped, layouts whose eager calls take at least as long are not captured.
• Timings are read from completed events on later calls, so no synchronization is added.

Results

TabArena, outer protocol, 816 splits, Kumo-Tabular-S (8 estimators, fp16, no KV cache), 8× NVIDIA RTX PRO 6000 Blackwell with one task per GPU; each branch was run on its own:

	#994	This PR
Elo	1793.3	1793.3
Median infer s/1K	0.0868	0.0804 (−7.4%)

• Predictions are identical to Make estimator batching exact, robust and automatic #994 on all 816 splits.
• Small tables that repeat across folds gain the most: 6 datasets run at least 1.2x faster, up to 1.47x (diabetes).
• Datasets with only 3 repeats and several batch shapes pay one capture per shape without enough replays to recover it (worst: seismic-bumps 0.86x, wine_quality 0.87x).
• On the largest tables (e.g., GiveMeSomeCredit, APSFailure, kddcup09_appetency), replays save at most 0.5% of the time while their pools would take 3–8% of device memory, so no graph is kept and no memory is held.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.