Peak-Memory Schedules You Can Trust: Exact Realization, Certified Optimality, and the Limits of Recomputation for Inference

Ivan Fomin; Ivan Koryabkin; Boris Sheludko

cs.LGOct 16, 2026

Peak-Memory Schedules You Can Trust: Exact Realization, Certified Optimality, and the Limits of Recomputation for Inference

Authors: Ivan Fomin, Ivan Koryabkin, Boris Sheludko

Organizations: Yandex, Moscow, Russia.

Abstract

Real-time inference on small devices needs low activation memory without extra latency, which makes operator order the most valuable memory lever: it is free in time. We show that the usual way of deploying a better order, handing it back to the compiler and regenerating code, is unreliable. In PyTorch Inductor the regenerated program misses its predicted peak in 22 of 54 CPU configurations. It is worse than the stock program in 9 of them and in 19 of 50 on a GPU, and on the GPU it also changes numerical results. We instead profile the generated program, model its memory exactly from its own statements and their read/write effects, and permute those statements. We prove that this realizes the modeled peak whenever allocation does not depend on history. It does so to the byte, with bitwise-identical results, in all our CPU runs. On the GPU we characterize why the caching allocator breaks the condition. To know when to stop, we give forward and reverse "trap" lower bounds, each computable with $O(n^2)$ minimum cuts. They certify the found order optimal for 94% of held-out wrapper graphs and all 16 microcontroller-class models we study, in under a second. So most inference graphs are already scheduled optimally. Branching networks are the exception, where reordering is free and pays, e.g. 10.0% on DeepLabV3. In compiled training we save up to 7.6%. Beyond the certified optimum we study selective recomputation with per-operator caps. It is strongly NP-hard for every fixed cap even at treewidth two, encodes knapsack exactly on series-parallel graphs, and is solvable exactly on bounded treewidth. It is useless on public tiny models whose peak is a single layer, but on architectures that hold a cheap value for a late reuse it pays: on the public ReDimNet speaker networks, whose order is already optimal, it lowers the peak by 5.2-6.1% at under 0.4% extra MACs, with bitwise-identical outputs.

Figures & tables

Panel of Figure 1 — Figure 1: Fixed-wrapper realization (blue) versus the usual regeneration pipeline (orange). The order is optimized on a model built from the generated program itself, and then applied to that program by permuting its statements, so the kernels, their arguments and buffer reuse are those that were profiled.

Figure 2: (a) Fraction of instances on which a lower bound equals the optimum. Blue: random DAGs with exactly computed optima. Orange: forward and backward wrapper graphs of the held-out training and inference models, where the reference is the peak of the found order. (b) Time to certify that

\mathrm{OPT}-1

is infeasible, trap certificates versus an exact CP-SAT cumulative model with one worker and a 120 s limit, on the certified graphs.

Figure 3: The same optimized order, realized by permuting the generated program (x) versus regenerating code for it (y). Points below the diagonal lose part of the saving and points below zero are worse than the stock program; one CPU point at

-30.6\%

is clipped.

Figure 4: Live activation memory along four microcontroller-class models, executed in model order, normalized by the peak. Each peak is set by a single operator (an early high-resolution layer, or a U-Net skip concatenation), which bounds what any schedule of whole tensors can achieve.

Size	Ops	GMACs	Opt. order	Recompute	MACs
B0	309	0.59	7.28	6.85 ( $-$ 6.0%)	$+$ 0.33%
B1	223	0.85	10.49	9.85 ( $-$ 6.1%)	$+$ 0.33%
B2	364	1.51	16.78	15.91 ( $-$ 5.2%)	$+$ 0.25%
B3	651	4.40	16.78	15.91 ( $-$ 5.2%)	$+$ 0.09%
B4	537	7.05	33.56	31.76 ( $-$ 5.4%)	$+$ 0.11%
B5	537	13.21	33.56	31.76 ( $-$ 5.4%)	$+$ 0.06%

Table 1: Selective recomputation on the public ReDimNet backbones (2 s input, fp32 activations). Peak live activation memory (MB) of the certified optimal order, before and after recomputing the stem for its late uses, and the extra MACs. All outputs are bitwise identical; both orders are certified optimal for their graphs.

Group	$n$	$>0$	Median %	Max %
Replication of earlier rows	9	7	2.54	7.60
Transformers (held-out)	15	14	1.30	2.75
CNNs (held-out)	8	3	0.00	3.02
BERT/GPT-2 batch $\times$ sequence	18	17	1.30	3.46

Table 2: Peak saving of fixed-wrapper realization over stock Inductor (memory pass on), compiled training on CPU. The prediction error is 0 B and gradients are bitwise identical in every run.

Figure 5: Live memory at each unit of the backward wrapper of a two-layer Transformer encoder. Deferring two parameter-gradient materializations past the flash-attention backward call lowers the peak to the certified lower bound.

Figure 6: Saving of the same reordering on CPU and on GPU (peak allocated by the CUDA caching allocator). Filled markers: the GPU peak equals the model prediction. Hollow squares: the allocator’s history-dependent block sizes break the prediction.

Restriction	Result
fixed uniform cap $k\geq 1$ , treewidth $\leq 2$	strongly NP-complete
caps in binary	PSPACE-complete (classical)
min. work, uniform $k{=}2$ , TTSP, out-degree $\leq 2$	equals 0–1 knapsack
caps $\leq 3$ , TTSP, out-degree $\leq 3$	NP-complete
(peak, work) frontier, $k{=}2$ , TTSP	$2^{m}$ points on $O(m)$ vertices
min. work, tree decomposition of width $t$	$f(t,K,B)\,\mathrm{poly}(n,L)$

Table 3: Complexity of peak feasibility with recomputation caps (minimum work where noted). TTSP: two-terminal series–parallel graphs.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Configuration	Params (M)	Stock Inductor (B)	Fixed wrapper (B)	Saving %	Regen. saving %	Cert. opt.
TransformerEnc 2L $d$ 64 $s$ 128	0.1	1,074,056	1,024,904	4.58	4.58	2/2
TransformerEnc 4L $d$ 64 $s$ 128	0.2	1,932,168	1,883,016	2.54	2.54	2/2
TransformerEnc 2L $d$ 256 $s$ 128	1.6	6,842,376	6,842,376	0.00	0.00	2/2
TransformerEnc 2L $d$ 256 $s$ 256	1.6	8,678,408	8,542,472	1.57	1.57	2/2
TransformerEnc 2L $d$ 256 $s$ 512	1.6	16,813,320	16,026,888	4.68	4.68	2/2
TransformerEnc 2L $d$ 256 $s$ 1024	1.6	32,557,064	31,520,776	3.18	3.18	2/2

Appendix

Table 4: Compiled training on CPU, main grid. Saving is relative to stock Inductor with its memory pass. The prediction error is the largest absolute error over the wrapper windows of the fixed-wrapper run. “Regen.” hands the same order back to Inductor. “Cert. opt.” counts wrapper graphs for which

\max(\mathcal{F},\mathcal{R})

equals the found peak.

Configuration	Units	Stock (MB)	Ours (MB)	Saving %	$\max(\mathcal{F},\mathcal{R})$ (MB)	Cert.	Bitwise
BERT-base, $b$ 8 $s$ 512	133	75.50	75.50	0.0	75.50	yes	yes
ConvNeXt-T, $b$ 8	115	58.00	58.00	0.0	58.00	yes	yes
DeepLabV3-R50, $b$ 2, 512 2	119	167.77	150.99	10.0	150.99	yes	yes
DenseNet-121, $b$ 8	242	144.51	144.51	0.0	144.51	yes	yes
DistilBERT, $b$ 8 $s$ 512	67	75.50	75.50	0.0	75.50	yes	yes
EfficientNet-B0, $b$ 8	164	48.17	48.17	0.0	48.17	yes	yes

Appendix

Table 5: Inference on CPU (fp32, evaluation mode). Peak of the stock and of the reordered program, the best lower bound, whether it certifies the order, and bitwise equality with the stock program.

Configuration	Stock (MB)	Ours (MB)	Saving %	Pred. err (B)	Regen. %	Det. stock	Bitwise
bert-base-b1-s128	438.79	438.79	0.00	0	-0.45	no	within
bert-base-b2-s256	468.94	457.02	2.54	11,667,456	2.10	no	within
convnext-tiny-b2	239.08	235.99	1.29	368,640	0.92	no	within
densenet121-b2	261.49	261.49	0.00	0	-2.59	no	within
distilbert-b1-s128	270.12	270.12	0.00	0	-0.73	yes	yes
distilbert-b2-s256	288.08	282.84	1.82	6,436,352	3.00	yes	yes

Appendix

Table 6: GPU replication on an RTX 2080 Ti. Peak allocated by the CUDA caching allocator inside the measured step.

Model	Ops	MMACs	Peak (KB)	Arena (KB)	Order opt.	Exact
FC autoencoder (AD)	11	0.3	0.8	0.8	yes	yes
EfficientNet-B0, 128 2	134	130.4	491.5	491.5	yes	yes
GoogLeNet, 96 2	82	277.1	184.3	184.3	yes	yes
ResNet-8 (CIFAR)	15	12.6	49.2	49.2	yes	yes
DS-CNN (KWS)	12	3.3	19.2	19.2	yes	yes
TC-ResNet8-1.5	16	3.4	6.5	6.5	yes	yes

Appendix

Table 7: Microcontroller-class models (batch 1, int8 activations). Peak live activations and greedy arena of the model order, whether the order is certified optimal, and whether executing the schedule is bitwise exact.

Size	Frames	Opt. peak	Recompute	Arena before	Arena after	MACs %
B0	201	7.28	6.85 ( $-$ 6.0%)	7.29	6.85	0.33
B0	81	2.96	2.79 ( $-$ 6.0%)	2.97	2.79	0.34
B1	201	10.49	9.85 ( $-$ 6.1%)	10.49	9.85	0.33
B1	81	4.27	4.01 ( $-$ 6.1%)	4.27	4.01	0.34
B2	201	16.78	15.91 ( $-$ 5.2%)	16.79	15.91	0.25
B2	81	6.83	6.47 ( $-$ 5.2%)	6.83	6.48	0.25

Appendix

Table 8: Public ReDimNet backbones: peak of the certified optimal order, after recomputing the stem (and the saving), greedy arena before and after, and extra MACs.

Model	Ops	GMACs	Optimal peak (MB)	Cert.	With recomputation (MB)
YOLOv5nu	69	3.9	11.47	yes	11.47 ( $-$ 0.0%)
YOLOv8n	65	4.4	11.47	yes	11.47 ( $-$ 0.0%)
YOLO11n	101	3.3	14.34	yes	14.34 ( $-$ 0.0%)
YOLOv8s	68	14.3	22.12	yes	21.71 ( $-$ 1.9%)
YOLO11s	103	10.8	28.67	yes	28.67 ( $-$ 0.0%)
YOLOv8m	84	39.5	39.32	yes	39.32 ( $-$ 0.0%)

Appendix

Table 9: YOLO detectors at

640\times 640

(fp32, raw head outputs). The model order is certified optimal in every case.

Family	$n$	Mandatory	$G$	$\mathcal{F}$	$\mathcal{R}$	$\max(\mathcal{F},\mathcal{R})$	Worst $\mathrm{OPT}/\max$
random	160	63	111	116	143	148	1.188
sp	160	137	145	146	156	157	1.231
pumpkin	160	75	113	116	145	148	1.231
intree	160	142	142	142	160	160	1.000
training	160	160	160	160	160	160	1.000

Appendix

Table 10: Number of random DAGs on which each lower bound equals the exact optimum, and the worst ratio of the optimum to

\max(\mathcal{F},\mathcal{R})

Graph	$K$	$B$	$n{=}6$	$n{=}8$	$n{=}10$	$n{=}12$	$n{=}14$
ladder	1	4	4	34	17	232	118
ladder	1	6	313	268	121	4,849	595
ladder	1	8	1,615	3,280	41,656	43,055	2,705
ladder	2	4	10	659	169	45,831	17,554
ladder	2	6	105,131	128,836	31,570	$>3{\cdot}10^{6}$	1,101,844
ladder	2	8	$>3{\cdot}10^{6}$	$>3{\cdot}10^{6}$	$>3{\cdot}10^{6}$	$>3{\cdot}10^{6}$	$>3{\cdot}10^{6}$

Appendix

Table 11: States generated by the complete-trace dynamic program (limit

3\cdot 10^{6}

) on ladder graphs (width 2) and two-arm pumpkins, with

K

the cap and

B

the memory budget.

Inference-Time Optimization

Inference-time optimization adapts model computation, outputs, or runtime settings to improve quality under deployment constraints.

See all 96 papers

Sep 24, 2026

KernelOPT: Dispatch-Aware Agentic Search for GPU Kernel Optimization

Deep learning inference and training performance depends critically on GPU kernel efficiency. Modern compilers such as PyTorch Inductor automatically generate GPU kernels from high-level model code, but frequently underperform expert-written implementations by wide margins. Recent LLM-assisted kernel optimizers can close this gap for standalone kernels, yet treat compiled models as black boxes, generally optimizing individual standalone kernels without respecting the compiler's structural decisions or verifying the model end-to-end. We present KernelOPT, a multi-agent system that treats compiled models as structured artifacts. It preserves vendor library calls (cuBLAS, cuDNN) and exclusively targets generated Triton sub-kernels using five profiling-guided LLM agents. A four-gate verification cascade applies static validation, multi-seed correctness checking, model-level float64-fallback verification, and performance gating (

γ{=}1.03

) to filter candidates and verify the re-stitched model end-to-end. When candidates fail verification, the system preserves the compiler baseline. The system accepts PyTorch nn Modules, standalone Triton kernels, and Helion kernels. Evaluated on 250 KernelBench problems (100 Level 1, 100 Level 2 and 50 Level 3) on NVIDIA H200, KernelOPT achieves geometric mean speedups over torch compile of 1.40

\times

(L1), 1.15

\times

(L2), and 1.07

\times

(L3) across all kernels, including fallback cases. Optimized-only geomeans (excluding cases where verification gates preserve the compiler baseline) are substantially higher: 2.54

\times

(L1: 36/100), 1.84

\times

(L2: 23/100), and 1.37

\times

(L3: 11/50), reflecting where the optimizer achieves meaningful leverage.

Aheli Poddar, Sanskar Prasad, Arindam Samanta +3

Red Hat

GPU Kernel Optimization Agentic Search Large Language Model-Guided Optimization Compiler Optimization Neural Network Hardware Acceleration

Ask PDF

Aug 21, 2026

Tensor Seeks Layout: Formalizing Layout Selection for ML Compilers

Modern machine learning compilers select tensor memory layouts to minimize execution cost under hardware constraints. Layout selection is global: an operator may be fastest under one layout while its consumers prefer another, and aligning these preferences requires explicit layout conversions that can hurt model performance. Despite its practical importance, layout selection lacks a formal basis, so current compilers rely on ad-hoc heuristics. This paper presents the first formal study of layout selection in machine learning compilers. We formulate the problem as combinatorial optimization over dataflow graphs, minimizing the sum of operator execution costs and the per-tensor cost of these conversions. Our theoretical analysis shows that optimal layout selection is computationally hard, even for programs containing only matrix multiplications over two-dimensional tensors. We design an optimal polynomial-time algorithm for dataflow graphs of bounded treewidth. For general instances, we give a weighted MaxSAT encoding that an off-the-shelf solver can optimize. The formulation unifies several existing layout optimization strategies, including XLA's layout assignment, partition dimension selection in systolic array compilers, and layout planning in mobile GPU optimizers. We implement the formalization in a production compiler for an AI accelerator and measure the execution time of the compiled models under greedy heuristics, the compiler's rule-based strategy, and an optimal solver. Simple heuristics degrade execution time by up to

5\times

on some workloads. Where the compiler's cost model is accurate, the solver matches or beats the rule-based strategy. On workloads with complex data movement it falls behind, and since the solver minimizes the stated objective exactly, that gap isolates cost-model error from search quality, showing where compiler effort actually pays off.

Clemens Eisenhofer, Yuwen Jia, Daniel Kroening +1

TU Wien, Vienna, Austria · Amazon, Cupertino, California, USA

Compiler Optimization Deep Learning Optimization Combinatorial Optimization

Ask PDF

Jul 18, 2026

Cost Accounting for Reactive Computational Graphs: Exhaustive Sweeps, Sequential Mutation, and the Backward-Locality Gap

Exhaustive site-by-site interventions on a neural network's computational graph -- activation-patching sweeps, circuit-discovery searches, systematic ablation studies -- mutate the graph at every candidate site, and their cost is dominated by recomputation after each mutation. On a reactive graph engine whose invalidation provably touches exactly the downstream cone of a mutated node, we give a complete cost accounting for such workloads. First, the aggregate speedup of an exhaustive sweep over independent full recomputations is not a universal constant: if per-layer weight varies regularly with depth at Karamata index q, the ratio converges to (q+2)/(q+1) when weight concentrates near the output and to q+2 near the input, recovering 2 only in the depth-uniform case; a wall-clock corollary predicts a ceiling of about 1.79, below 2, until interpreter overhead is compiled away. Second, we prove the exact cost of a sequence of persistent mutations, never undone between insertions: the interleaved cost exceeds the isolated sum by an exact overcount summed over comparable site pairs, with closed-form extremes over insertion orders, while batched application is order-independent and sub-additive, costing exactly the union of the sites' cones plus the fresh nodes. Third, we prove the exact mirror of forward locality for the backward pass, showing it collapses the aggregate speedup to 1 under backpropagation on architectures without long skip connections. Every identity is validated on NeuroDSL, a reactive graph engine in Julia: measured sweep ratios converge to the predicted limits under four cost profiles; the training-mode ratio collapses to 1 at the predicted rate; and all 18 per-graft sequential costs and the batched total match the closed forms at zero tolerance across three insertion orders.

Abdallah Khemais

ISITCOM, University of Sousse

Activation Patching

Ask PDF

Abstract

Inference-Time Optimization

Explore similar work

KernelOPT: Dispatch-Aware Agentic Search for GPU Kernel Optimization

Tensor Seeks Layout: Formalizing Layout Selection for ML Compilers

Cost Accounting for Reactive Computational Graphs: Exhaustive Sweeps, Sequential Mutation, and the Backward-Locality Gap