Static Memory Planning from Tensor Liveness Intervals

Once the graph has been fused and scheduled, every remaining intermediate has a known size and a known position in the execution order. That is enough to answer the only question the runtime cares about: where in device memory does each tensor live. Allocating each one separately from a general-purpose allocator is the obvious answer and the wrong one. A caching allocator makes allocation fast but cannot prevent fragmentation across repeated inference runs with varying shapes, and the fragmentation shows up as an out-of-memory failure at a working set that ought to fit.

The alternative is to allocate one large region at load time and hand out fixed offsets within it. The plan is computed once, from the graph, and it never changes. This is the same problem a register allocator solves, with a large address space instead of sixteen registers, and it inherits the same theory.

Deriving the interval

A tensor’s lifetime starts at the node that writes it and ends at its last reader. If the tensor is a graph input, its lifetime starts at the beginning; if it is a graph output, it ends at the end. Converting nodes to positions in the topological order gives each tensor a half-open interval $[s, e)$ over the schedule.

The endpoint rule matters. If tensor A is written at position 3 and last read at position 7, and tensor B is written at position 7, the two may share storage: at the moment B is produced, A is no longer needed. Treating the last read as inclusive and the first write as inclusive would forbid that sharing and waste the buffer. The rule that works is that a tensor is live at position $p$ if $p$ is at or after its definition and before its last use, and a consumer that reads both a dying tensor and a newborn one needs both simultaneously, so the dying tensor’s interval must extend through the consumer.

Fusion changes these intervals directly. Fusing a producer into its consumer deletes the producer’s output tensor entirely, along with its interval, which is the main reason a fused graph plans into less memory than an unfused one. Conversely, choosing a topological order that emits a branch early extends that branch’s interval across the nodes that follow. Liveness is not a property of the graph; it is a property of the schedule.

def liveness(nodes, outputs):
    """Half-open [start, end) intervals over topological position."""
    start, end = {}, {}
    for pos, node in enumerate(nodes):
        for v in node.inputs:
            if v not in start:          # graph input: born at time 0
                start[v] = 0
            end[v] = pos + 1            # extend through this consumer
        for v in node.outputs:
            start[v] = pos
            end.setdefault(v, pos)
    for v in outputs:
        end[v] = len(nodes)
    return {v: (start[v], end[v]) for v in start}

Interval colouring, and why it is not general graph colouring

Two tensors may share an offset if and only if their intervals do not overlap. The relation “overlaps” on intervals is an interval graph, and colouring an interval graph is solvable in polynomial time: sort by start, and assign each interval the lowest-numbered colour not used by any interval that overlaps it and starts before it ends. The greedy sweep is optimal for the number of colours, a result that follows from the structure of interval graphs being chordal and perfect.

That optimality is why memory planning is easier than register allocation for a general program. In a compiler with branches, two values defined on different paths may never be simultaneously live but a naive liveness analysis marks them as overlapping, and the interference graph is not an interval graph. In a straight-line dataflow graph with a fixed schedule, it is. The caveat is that the colours here are not unit-sized registers: tensors have different byte sizes, so “colour” becomes “offset”, and the problem becomes a variant of strip packing rather than pure colouring.

The standard reduction handles this by splitting: the largest tensor’s size defines a granularity, and every other tensor is rounded up to a multiple of it, at which point the problem is colouring with weighted colours. The rounding is where the bytes go. A 3 MB tensor and a 200 KB tensor both occupy 3 MB in the plan if the granularity is 3 MB, so a graph with many small intermediates and one large one wastes most of the arena. Compilers reduce the loss by planning in two tiers: large tensors get offsets from a size-ordered pass, and small tensors are packed into the gaps left behind.

def plan_offsets(intervals, sizes, align=256):
    """Greedy best-fit: place each tensor at the lowest free offset."""
    placed = []                              # (start, end, offset, size)
    offsets = {}
    for v in sorted(intervals, key=lambda t: -sizes[t]):
        lo, hi = intervals[v]
        size = (sizes[v] + align - 1) // align * align
        busy = sorted((o, o + sz) for st, en, o, sz in placed if st < hi and en > lo)
        cursor = 0
        for b0, b1 in busy:                  # find first gap of `size` bytes
            if b0 - cursor >= size:
                break
            cursor = max(cursor, b1)
        offsets[v] = cursor
        placed.append((lo, hi, cursor, size))
    return offsets

Sorting by decreasing size before placing is the heuristic that makes greedy packing behave well, and it is the same ordering used in offline strip packing. It is not optimal. Constructed instances exist where best-fit-decreasing wastes a constant fraction of the arena, and a graph whose intervals nest badly will find them. The offset is bounded in practice because inference graphs have a repetitive structure: the same block repeated many times produces the same liveness pattern each time, so a planner that handles one block well handles all of them.

In-place reuse and the aliasing hazard

The plan assumes a tensor’s storage may be reused as soon as its interval ends. That assumption is only sound if the write to the reused region happens after the last read of the dying tensor, which the schedule guarantees, and if nothing else holds a reference to the dying buffer. A graph output held by the caller after the call returns must be excluded from reuse, or the caller reads overwritten data. The same applies to any tensor whose address has escaped to an external buffer, which is why view-like operators need explicit treatment: a reshape recorded as a view means two values share storage, and the planner must either merge their intervals or force a copy.

Compilers express this by computing a set of tensors that must be materialised into caller-owned memory and treating their intervals as extending past the end of the graph. The cost is real: a model with many outputs cannot reuse their buffers for later computation, and its arena is correspondingly larger.

The second hazard is a mismatch between the plan and the actual runtime. A plan computed for a static shape is exact for that shape. If a request arrives with a different batch size, the offsets remain valid only if the arena was sized for the maximum. Compilers that support a range of shapes plan for the worst case, which means the arena is sized by the largest batch and every smaller request wastes the difference. Planning per shape and caching the plans recovers the waste at the cost of a compilation per distinct shape, the trade-off that shape specialisation is built around.

What the arena does not solve

A single arena removes fragmentation within one graph execution. It does not remove the need for multiple arenas when the device has distinct memory spaces, such as shared memory per block, which is not planned statically by the graph compiler at all but by the kernel’s tiling. It also does not help with the memory that must persist across executions, namely the weights, which are constant and are typically placed in a separate read-only region so that they can be shared between concurrent executions of the same model.

Arena planning is the point at which the compiler’s decisions become visible as a memory number. A graph with $N$ intermediates, each of size $S$, and no reuse needs $NS$ bytes; a good plan with a schedule that keeps the live set to $k$ tensors needs roughly $kS$. On a transformer block, $k$ is small relative to $N$ because the block is a chain, so the saving is large. On a graph with many parallel branches that all feed a final concatenation, $k$ approaches $N$ and planning saves little, because the concatenation genuinely needs every branch’s output alive at once. That is a property of the model’s structure, not of the planner, and it is the reason fusion into the concatenation is worth attempting even when the concatenation is large.