TelaMalloc: Efficient On-Chip Memory Allocation for Production Machine Learning Accelerators

Before an ML model can run on a chip like a TPU or a Pixel phone, the compiler must fit every intermediate tensor into a small fixed on-chip memory, choosing an address for each so that at no moment do the live tensors overflow — a 2D rectangle-packing puzzle (each buffer is a rectangle in time-versus-memory) that is NP-hard, and the two known tools each fail — fast greedy heuristics get stuck or waste up to 43% extra memory, while exact constraint solvers are correct but so slow they were 68% of one model's compile time and could take minutes (undeployable on a phone); TelaMalloc (Google) rejects that either-or by fusing them — a domain-aware heuristic places most buffers quickly, and a constraint solver runs alongside to catch the instant a choice makes the packing impossible and to guide the backtracking — matching heuristic speed on easy cases while running 2-100x faster than the solver on hard ones, with 27-37x fewer failures, now shipping in the Pixel 6 and Google's datacenters.

The argument in one glance

Problem

Every tensor needs a home in a tiny on-chip memory, and they must not collide

An ML accelerator — a datacenter TPU, or the Tensor chip in a Pixel phone — runs its computation out of a small, fast on-chip SRAM scratchpad. Before the model can run, the compiler must give every intermediate tensor (buffer) a specific ADDRESS in that memory, and it must do so such that at every moment during execution the buffers that are simultaneously alive fit within the chip's fixed capacity without overlapping. Because the computation graph is known ahead of time, each buffer's LIFETIME (when it is created and last used) is fixed; what the compiler chooses is WHERE each buffer sits. Get it wrong and either the model does not fit and cannot run at all, or it wastes precious SRAM. This allocation is a genuine bottleneck: on phones it runs on-device at launch, and in datacenters it runs over and over inside compiler optimization loops.
What people did

Use a fast greedy allocator, or an exact solver

There are two well-worn tools. A greedy HEURISTIC allocator — like TensorFlow's BFC or an XLA skyline packer — places buffers one at a time using simple rules (biggest first, stack on top of what is there), which is fast. Or an exact SOLVER — integer linear programming or constraint programming — that models the whole packing as a set of constraints and finds a provably valid (or optimal) placement. Fast-and-approximate or slow-and-exact: the classic pair of options for an NP-hard packing problem, and prior systems picked one.
Why it fell short

Heuristics get stuck; solvers are too slow and wildly unpredictable

Each option fails at production scale. HEURISTICS are fast but brittle: on hard instances they paint themselves into a corner, either failing to find any valid packing or needing far more memory than necessary — one case requires 43% more memory than optimal — and a failed allocation means the model simply does not run. SOLVERS are correct but do not scale: the ILP allocator previously used at Google was 68% of total compilation time for one real model, and some models took tens of seconds to minutes, which is impossible for on-device compilation on a phone. Worse, solver runtime is wildly unpredictable — changing the memory size by a single byte can send the solver down a completely different search path, and one model took over 9 hours in some runs. So one tool is fast but unreliable, and the other reliable but unusably slow and erratic; neither alone works in production.
This paper's idea

Fuse the heuristic and the solver into one search

TelaMalloc (Google — Maas, Beaugnon, Chauhan, Ilbeyi) refuses the either-or and couples the two approaches into a single search. A domain-aware HEURISTIC drives: at each step it picks which buffer to place next and where, using rules like longest-lifetime-first, so the common case is as fast as a heuristic. But a constraint SOLVER runs alongside it, and after each placement it checks whether that choice just made the remaining problem UNSOLVABLE — pruning dead-ends the instant they appear rather than after wasted work — and when the search must back up, the solver's conflict information directs it to jump straight back to the real source of the conflict. A preprocessing step solves the most memory-contended time phases first, in isolation. The result matches heuristic speed on easy instances and runs 2-100x faster than the pure solver on hard ones, with 27-37x fewer failures — now shipping in the Pixel 6 and in Google's TPU compiler.

Why this problem is worth solving

Memory allocation on an ML accelerator decides whether a model runs at all and how well, and because the accelerator’s on-chip memory is small and fixed, the packing has to be tight and correct — which is exactly what makes it hard. The task is best seen as a

the geometric view of on-chip allocation. Draw a strip whose width is time (the execution schedule) and whose height is the on-chip memory capacity. Each buffer is a RECTANGLE: its horizontal extent is its lifetime (from when the graph creates it to its last use), which is FIXED by the computation graph, and its height is its size. The compiler must choose each rectangle’s VERTICAL position (its memory offset) so that no two rectangles overlap and none pokes above the top of the strip — because two overlapping rectangles mean two buffers assigned the same memory at the same time, and poking above the strip means exceeding capacity. Packing rectangles into a fixed-height strip without overlap is a classic NP-hard problem, and this is that problem with the horizontal positions pinned and only the vertical positions free

. Since the goal is to use no more than the available memory, the allocator is not just optimizing quality — on a tight model it must find a valid packing at all, or the model cannot launch.

What separates this from ordinary memory allocation, and what makes a heavyweight approach even thinkable, is that the whole problem is known in advance. This is

the crucial difference from dynamic memory allocation in a language like C++, where allocation and free requests arrive one at a time, unpredictably, and the allocator must respond immediately with no knowledge of what comes next. An ML computation graph, by contrast, is STATIC and known before execution begins: the compiler knows every buffer that will ever be needed, exactly when each is born and dies, and how big each is. This changes the problem from an online, react-as-you-go one into an offline, plan-the-whole-thing one — which means the allocator can consider all the buffers together and globally optimize the packing, rather than greedily reacting to each request as it arrives

. That offline knowledge is what makes it worth running a solver at all — you can afford to think hard because you are compiling once and the layout is reused for every inference. But it also raises the bar: because you COULD find a near-optimal packing, wasting memory or failing is a real, avoidable loss. The whole challenge is to get the quality and reliability that the static, known-ahead problem makes possible, at a speed that on-device compilation and datacenter compile loops can actually afford — which is what the heuristic-alone and solver-alone approaches each fail to do.

What people did before, and where it broke

Greedy heuristic allocators (TensorFlow BFC, XLA skyline). Place buffers one at a time by simple rules (largest first, stack on the skyline). Fast and simple, but on hard instances they get stuck — failing to find a valid packing, or using far more memory than needed (43% over optimal in one case) — and a failure means the model does not run.
Exact ILP / constraint solvers. Model the full packing and solve for a valid or optimal layout. Provably correct, but do not scale: 68% of compile time for one model, minutes on large graphs, undeployable on a phone — and with pathological unpredictability, where a one-byte change in memory size can swing runtime to over 9 hours.
Heuristic-then-solver fallback (run them in sequence). Try the heuristic first and fall back to the solver only if it fails. Better than either alone, but it still pays the full solver cost on every hard instance, and the two never share information — the heuristic cannot use what the solver learns, or vice versa.
Optimizing quality without a compile-time budget. Prior work aimed for the best packing regardless of allocation time, fine for offline research but unusable where the allocation is on the launch critical path of an app or inside a tight compile loop.
The gap. No allocator combined heuristic SPEED with solver RELIABILITY in a single, tightly-coupled search — heuristic decisions guided and pruned by a live solver — so as to be fast on easy cases and reliable on hard ones within a production budget. TelaMalloc does exactly this.

How TelaMalloc works

Drive with a heuristic, check and guide with a solver, and solve the hard phases first

Frame allocation as rectangle packing

Each buffer is a rectangle whose time extent is fixed by the graph and whose memory offset is free; TelaMalloc must place them all in the fixed-height memory strip with no overlaps. Non-overlap is encoded as a constraint for every pair of time-overlapping buffers: one must sit entirely above the other in memory.
↓ →
Reject the fast-versus-correct dichotomy new

Instead of a fast-but-brittle heuristic OR a correct-but-slow solver, TelaMalloc runs both in one loop: a heuristic proposes the next placement, and a constraint solver immediately validates it and feeds back. The heuristic supplies speed and domain knowledge; the solver supplies the guarantee that a decision has not doomed the rest.
↓ →
Place buffers with a domain-aware heuristic new

At each step the heuristic picks which buffer to place (trying longest-lifetime first, then largest size, then largest area) and where — either on the skyline of already-placed buffers (Tetris-style), or, better, by asking the solver for the minimum valid address so a block can pack UNDERNEATH earlier ones instead of only stacking on top.
↓ →
Let the solver prune dead-ends the instant they appear new

After each placement the constraint solver checks whether the remaining problem is still solvable. If a choice has just made a valid packing impossible, the solver catches it immediately and the branch is abandoned — so the search never wastes effort exploring a subtree that can never lead to a solution, which is exactly where a pure heuristic gets stuck.
↓ →
Backtrack toward the real conflict new

When all candidates at a decision point fail, TelaMalloc does not backtrack a fixed number of steps; it uses the solver's CONFLICT information to jump directly back to the placement actually responsible for the dead-end (the second-to-last conflicting decision). This conflict-directed backjumping skips over regions of the search space that could never have worked.
↓ →
Solve the most contended time phases first new

A preprocessing step identifies the time windows where many buffers are alive at once and memory is most contended. TelaMalloc solves those high-contention phases first, mostly in isolation — pinning down the hardest decisions early, when they most constrain everything else, which sharply reduces the branching for the rest of the packing.

Evaluated on production ML models across 1,192 input configurations, deployed in the Google Pixel 6 (on-device compilation) and the XLA compiler for TPUv4. An experimental learned (gradient-boosted-tree) backtracking policy, trained on 301,321 examples, further helps the long-tail hardest models.

The idea

Reject the false choice by making the two tools cover each other’s weakness. The framing everyone inherited was a dichotomy: a heuristic is fast but gets stuck, a solver is reliable but too slow, pick one. TelaMalloc’s central move is to see that these are not competitors but complements — the heuristic is weak exactly where the solver is strong (recognizing that a partial packing has become hopeless) and slow exactly where the heuristic is fast (making the many easy local decisions). So it fuses them into one search where the heuristic makes the bulk of the decisions at heuristic speed and the solver, running alongside, catches the moment a decision dooms the rest. This is the combine-the-approaches- that-each-fail-alone discipline the site keeps rewarding — the FHE-compiler page pairing simulated annealing with an ILP for alignment, the compiler-autotuning page wrapping a surrogate model in an acquisition search — where the win is not a cleverer single method but the recognition that two methods’ failure modes are opposite, so coupling them cancels both.

Use the expensive checker to make the cheap search smart. The solver in TelaMalloc is not a fallback that runs when the heuristic gives up; it runs continuously as a guide. It PRUNES — the instant a placement makes the remaining problem unsolvable, the branch is abandoned, so the heuristic never wastes time in a doomed subtree. And it DIRECTS backtracking — when the search hits a wall, the solver’s conflict information says which earlier decision actually caused the wall, so the search jumps straight there instead of undoing arbitrary recent steps. This is the same guide-the-search- with-a-model discipline the site sees when a learned cost model tells an autotuner where to look, or a monotonicity property lets a collectives search prune provably — the expensive component earns its keep not by solving the whole problem but by keeping the cheap search out of the parts that cannot pay off. Conflict-directed backjumping is the classic version of this from SAT solving, applied here to memory packing.

Attack the hardest subproblem first, while you still can. The order in which you make decisions in a constrained search matters enormously, because early decisions in the wrong place force endless backtracking later. TelaMalloc preprocesses to find the time phases where memory is most contended — where the packing is tightest and the constraints bite hardest — and solves those first, in isolation. Pinning down the hardest decisions early, when the most freedom remains to satisfy them, shrinks the branching of everything that follows. This is the tackle-the-bottleneck-first discipline the site rewards — the heterogeneous-serving page placing the highest-traffic stages first, the taint-analysis page starting the trace where it is tractable — the general principle that when part of a problem is much harder than the rest, you should confront it while your options are widest, not after the easy decisions have boxed you in.

Packing tensors into on-chip memory without the solver stalling

Where each tool alone fails, and how the fusion wins

Allocating a static ML graph's buffers into fixed on-chip SRAM. Numbers from TelaMalloc's evaluation.

Step 1 of 5

The packing puzzle

The compiler must place every tensor rectangle into the memory strip so that, at each moment, the live buffers fit within capacity without overlapping. The time extents are fixed by the graph; only the memory offsets are free. It is NP-hard, and on a tight model the allocator must find a valid packing or the model cannot run at all.

the graph is static, so the whole set of buffers and their lifetimes is known before execution
for every pair of time-overlapping buffers, one must sit entirely above the other in memory

Why it works

A heuristic gets stuck exactly where a solver is strong, and a solver is slow exactly where a heuristic is fast → fusing them lets the heuristic make the many easy decisions and the solver catch the few fatal ones → heuristic speed on easy instances, solver reliability on hard ones.
A pure heuristic cannot tell an early placement has doomed the rest until too late → the solver checks feasibility after each step and prunes dead-ends immediately → the search never wastes effort in a subtree that cannot lead to a solution.
Undoing arbitrary recent decisions after a dead-end is wasteful → the solver’s conflict information directs backjumping to the decision that actually caused it → the search skips regions that could never have worked.
Early decisions in the tightest region force the most backtracking → solving the highest-contention time phases first, in isolation, pins the hardest choices while freedom is widest → branching for the rest of the packing collapses.
The ML graph is static, so the whole packing is known in advance → TelaMalloc can globally optimize offline rather than react greedily → it affords a solver-grade search at a heuristic-grade speed, deployable on-device and in compile loops.

What it costs

Heavy overlap still stresses the solver. When many buffers are alive at once, the pairwise non-overlap constraints grow quadratically and up to 87% of the time is spent in the solver — 1,000 fully-overlapping buffers make each step exceed 100 milliseconds — a bottleneck inherited from the solver’s scaling.
A long tail remains unsolved. Even with the fusion, a small number of models cannot be allocated by any approach, which is what motivated the experimental learned-backtracking extension.
The learned policy is imperfect. The gradient-boosted-tree backtracking model occasionally misfires — 4 of 1,192 cases time out and 9 regress significantly — so it needs heuristic fallbacks, and a static offline model cannot adapt to new workloads after deployment.
Reproducibility constrains the design. Production requires deterministic, reproducible compilation, which rules out continuous online learning that could otherwise adapt the policy — a real constraint the abstract problem does not impose.
Extreme unpredictability is tamed, not eliminated. The fusion greatly reduces the solver’s pathological runtime swings, but the underlying sensitivity (a one-byte change altering the search path) is a property of constraint search that the design manages rather than removes.

The evaluation

Production ML models across 1,192 input configurations, comparing TelaMalloc against greedy heuristic allocators and a pure ILP solver on allocation time, failure rate, and steps to completion, plus end-to-end deployment in the Google Pixel 6 (on-device compilation) and the XLA compiler for TPUv4.

Allocating on-chip ML buffers

Property	Greedy heuristic (BFC / skyline)	Exact solver (ILP / CP)	TelaMalloc (hybrid)
Speed
Finds a valid packing?
Memory quality
Predictable compile time
Deployable on-device

TelaMalloc gets the speed of a heuristic and the reliability of a solver at once by coupling them, so it is fast on the easy models and succeeds on the hard ones — 27-37x fewer failures than heuristics, 2-100x faster than the solver — which is what makes on-device compilation on the Pixel 6 feasible and improves TPUv4 throughput by fitting more buffers on-chip.Underlined values explain themselves when clicked.

Where this stops working

What to keep in mind

TelaMalloc tames the solver’s cost but does not abolish it: when buffers heavily overlap, the pairwise non-overlap constraints grow quadratically and the solver dominates (up to 87% of the time, over 100 milliseconds per step at 1,000 fully-overlapping buffers), and a small long tail of models still cannot be allocated by any approach, which is why the authors add an experimental learned-backtracking policy — itself imperfect, occasionally timing out or regressing and so needing heuristic fallbacks. Production also imposes constraints the abstract problem does not: compilation must be deterministic and reproducible, ruling out the continuous online learning that might otherwise adapt the policy, and the underlying sensitivity of constraint search (a one-byte memory change altering the whole path) is managed rather than removed. None of this diminishes the core lesson — that the fast-heuristic-versus-slow-solver choice is a false dichotomy, because their weaknesses are opposite: a heuristic is weak precisely where a solver is strong (proving a partial packing hopeless) and slow precisely where a heuristic is fast (the many easy local decisions), so fusing them into one search — heuristic driving, solver pruning and directing backtracking, hardest phases solved first — yields heuristic speed with solver reliability, enough to ship on-device ML compilation on a phone and pack more onto a TPU.

Background

Fuse the two tools whose weaknesses are opposite, and let the checker guide the search

TelaMalloc sits in the site’s compilers-and-accelerators cluster and rhymes with several ideas covered elsewhere on this site. Its governing move is to reject a false dichotomy by combining two approaches whose failure modes are opposite: a fast heuristic that gets stuck and a slow solver that always succeeds, fused so the heuristic supplies speed on the easy decisions and the solver supplies reliability on the hard ones — the same combine-the-complementary- methods discipline the FHE-compiler page shows when it wraps simulated annealing around an ILP for alignment and the compiler-autotuning page shows when it pairs a surrogate model with an acquisition search, where the win comes not from a cleverer single method but from recognizing that two methods cancel each other’s weakness. Its second move is to let the expensive checker GUIDE the cheap search rather than merely back it up — the solver prunes doomed branches the instant they appear and its conflict information directs backjumping to the true cause — the guide-the-search-with-a-model instinct the autotuning page uses when a learned cost model steers where to look and the PCIe-collectives page uses when a monotonicity property makes pruning provable, with conflict-directed backjumping lifted straight from SAT solving. Its third move is to attack the hardest subproblem first — solving the most memory-contended time phases while freedom is widest — the tackle-the-bottleneck-first discipline the heterogeneous-serving page uses placing high-traffic stages first and the firmware-taint page uses starting where the trace is tractable. And underneath it all is an exploit-the-static-graph insight — because the whole allocation is known in advance, unlike dynamic malloc, it can be globally optimized offline, the same solve-once-because-you-know-the-future economics the heterogeneous-serving and autotuning pages rely on — plus the hard-won credibility of a system, like the datacenter-IO-control page, whose design was shaped by real production constraints (an on-device compile budget, reproducibility) rather than the abstract problem alone.

If you remember only three things

1. Running an ML model on an accelerator requires the compiler to pack every intermediate tensor into a small fixed on-chip memory — a 2D rectangle-packing problem (each buffer a rectangle fixed in time by the graph, free to place in memory) that is NP-hard and on the critical path (on-device on phones, in compile loops in datacenters); the two standard tools each fail, as greedy heuristics are fast but get stuck or waste up to 43% extra memory (and a failed allocation means the model will not run), while exact ILP/constraint solvers are reliable but far too slow (68% of one model's compile time, minutes on large graphs) and pathologically unpredictable (a one-byte change can mean a 9-hour compile).
2. TelaMalloc rejects this false dichotomy by fusing the two into one search: a domain-aware heuristic drives (placing buffers by longest-lifetime, largest-size rules, at heuristic speed), while a constraint solver runs alongside to prune any placement that makes the remaining packing unsolvable the instant it happens, and to direct conflict-based backjumping to the decision that actually caused a dead-end — plus a preprocessing step that solves the most memory-contended time phases first, in isolation, pinning the hardest choices while freedom is widest.
3. The fusion matches heuristic speed on easy instances and runs 2-100x faster than the pure solver on hard ones (median 4.7x on Pixel 6), with 27-37x fewer failures, and is production-deployed — shipping in the Google Pixel 6 (compiled on-device) and in the XLA compiler for TPUv4, where tighter packing yields up to 7% better end-to-end performance — with honest limits (heavy buffer overlap still stresses the quadratic-scaling solver, a long tail remains hard, and production reproducibility rules out online learning). The durable lessons: when two tools fail in opposite ways, fuse them so each covers the other; let the expensive checker guide the cheap search rather than just back it up; and attack the hardest subproblem while your options are widest.

Check yourself

Check your understanding

1. TelaMalloc frames the classic "fast heuristic versus slow solver" choice as a false dichotomy. Predict why fusing the two beats picking either.
2. In TelaMalloc the solver does not just run as a fallback when the heuristic fails; it runs continuously alongside the heuristic. Predict what the solver contributes by running throughout the search.
3. TelaMalloc preprocesses to find the time phases with the highest memory contention and solves those first, in isolation. Predict why decision ORDER matters so much in this constrained search.
4. The paper stresses that ML computation graphs are STATIC — the whole set of buffers and their lifetimes is known before execution — unlike dynamic malloc where requests arrive one at a time. Predict why this distinction is what makes TelaMalloc's approach possible.

Sources

ACM DL: TelaMalloc: Efficient On-Chip Memory Allocation for Production Machine Learning Accelerators — ASPLOS 2023 (full text read via the ACM PDF; no arXiv preprint).