Dynamic Shapes, Guards, and the Cost of Recompilation

Every fusion decision and every memory offset depends on extents being known. A tile of $128 \times 128$ over a matrix of $4096 \times 4096$ divides evenly; over a matrix of $4100 \times 4096$ it does not, and the kernel needs a masked tail. A memory plan computed for a batch of eight has offsets that are wrong for a batch of nine. The compiler therefore faces a choice it cannot avoid: either it compiles for one concrete shape, or it compiles something general enough to accept a family of shapes and pays for that generality in the kernel.

Serving systems encounter both regimes. An image classifier receives one image per request and always at the same resolution, so its shapes are fixed and the question does not arise. A language model receives sequences whose input length varies from a handful of tokens to tens of thousands, and whose decode step always processes one token per sequence but a variable number of sequences. The same compiled graph cannot be optimal for both.

What generality costs in the kernel

A kernel written for symbolic extents cannot use the tricks that make a specialised kernel fast. It cannot unroll a loop whose trip count is unknown, so it emits a counted loop with a branch. It cannot prove that a tile divides the extent, so it masks every load. It cannot size shared memory from a compile-time constant, so it either allocates the maximum or asks the driver at launch. Each of these is a small cost and together they are a substantial one: masked loads prevent the compiler from merging accesses into wide vector loads, and a dynamic loop bound prevents the scheduler from software-pipelining the iterations.

The magnitude of the cost is not a fixed ratio. It depends on how much of the kernel’s work sits in the tail, and it interacts with fusion, since a fused group’s tail is the tail of the longest-lived operand. A matrix multiply with extents that are large multiples of the tile size loses almost nothing to masking, because the masked path executes only for boundary tiles. A elementwise kernel over a dynamically sized vector loses more, because its whole body is a strided loop with a runtime bound and no vectorisation opportunity unless the compiler can prove alignment.

The pragmatic resolution used by most graph compilers is to keep the graph dynamic and specialise the kernel. The graph’s shape inference produces symbolic expressions, extents that are affine functions of a small number of symbols, and the kernel generator emits code parameterised by those symbols rather than by concrete values. This keeps one compiled artifact for the whole model while allowing the generator to emit a small dispatch on alignment, choosing a vectorised path when the pointer and extent permit and a scalar path otherwise.

Guards: the mechanism that makes specialisation safe

A specialised artifact is only valid for the shapes it was built for. Using it for a different shape either produces wrong results or corrupts memory, so the runtime must check before it dispatches. That check is a guard: a predicate over the properties the compilation assumed, evaluated on every call.

A guard records more than extents. It records dtype, device, the strides of each input, whether an input is contiguous, whether any tensor aliases another, and in some systems whether a particular value is a constant. The compiled code is dispatched only when all of them hold; otherwise the runtime takes a fallback path, which is either a generic compiled variant or an eager interpreter.

class Guard:
    """Predicate over the properties a specialised artifact assumed."""
    def __init__(self, shapes, dtypes, strides):
        self.shapes, self.dtypes, self.strides = shapes, dtypes, strides

    def holds(self, args):
        for a, shp, dt, st in zip(args, self.shapes, self.dtypes, self.strides):
            if a.dtype != dt or tuple(a.shape) != tuple(shp):
                return False
            # A stride mismatch means the tensor is not the layout we compiled for.
            if tuple(a.stride()) != tuple(st):
                return False
        return True

# Dispatch: one specialised artifact per distinct guard, keyed by the guard.
def lookup(cache, args):
    key = (tuple(a.shape for a in args),
           tuple(a.dtype for a in args),
           tuple(tuple(a.stride()) for a in args))
    art = cache.get(key)
    if art is None:
        return None            # miss: the caller compiles and stores a new one
    return art if art.guard.holds(args) else None

Two failure modes follow directly from that structure. The first is guard explosion: a model called with gradually increasing sequence lengths produces a new artifact per distinct length, and each compilation costs seconds. The cache fills with hundreds of artifacts, each holding its own weights-derived constants and its own memory plan, and the process runs out of memory compiling rather than inferring. The second is guard churn: a model whose inputs alternate between two shapes pays a compile on every other call, which is worse than never compiling.

Bucketing and padding as the alternative

The standard remedy for both is to quantise the shape space. Instead of compiling for every length, the compiler compiles for a small set of bucket sizes, and the runtime pads each request up to the next bucket. A sequence of 137 tokens runs as a 192-token bucket; a batch of 5 runs as a batch of 8. The number of artifacts is bounded by the number of buckets, and the guard only needs to check which bucket applies.

The cost is the padding itself. A 137-token request in a 192-token bucket processes 40 percent more token positions than the request contains, and attention work grows with the square of the sequence length, so the waste compounds: padding from 137 to 192 multiplies the score matrix from $137^2$ to $192^2$ elements, a factor of nearly two. For decode steps, where the per-token cost is dominated by reading the weight matrices rather than by attention, the waste is smaller and bucketing is nearly free. The choice of bucket boundaries is therefore workload-dependent: narrow buckets near the lengths that dominate the traffic distribution, wide buckets elsewhere.

Padding has a second, subtler cost. A padded token position still participates in the attention softmax unless it is masked, and a mask that is correct for one bucket is not correct for another. The kernel must be told the true length separately from the padded extent, which is an extra scalar argument and a branch in the innermost loop. That branch is cheap on a masked-load architecture and expensive on one where it breaks vectorisation, which is one of the reasons attention kernels are written with the mask applied to the score before the softmax rather than to the values after it.

Recompilation is a latency event, not a throughput event

When a guard fails and no cached artifact matches, the runtime compiles. That compilation happens on the critical path of the request that triggered it. For a large graph, the sequence of shape inference, fusion, layout assignment, memory planning, and code generation can take seconds, and the first request after a model load pays it while a client waits.

Systems reduce the exposure in three ways. They compile at load time for the shapes the deployment expects, so the common case never compiles. They compile asynchronously, returning results from a fallback path while a background compilation proceeds, so the triggering request is slow but not blocked. And they cache artifacts on disk keyed by the guard, so a restart or a new replica reuses the compilation rather than repeating it.

The third mechanism introduces a correctness obligation: the cache key must include everything the compilation depended on, including the compiler version, the device, and any flags. A key that omits the device produces an artifact whose layout choices were made for a different architecture, and the mismatch surfaces as a wrong result or a driver fault rather than as a cache miss.

The steady state is a distribution, not a shape

A serving system in operation does not see one shape; it sees a distribution over shapes, and the compiled artifacts it holds are a sample from that distribution chosen by the bucketing policy. The relevant question is not “what shape is the model compiled for” but “what fraction of requests land in a bucket whose padding is small”. A system that buckets at powers of two is efficient for requests near the top of each bucket and wasteful for those near the bottom, and the traffic distribution decides which of those dominates.

This is why the shape question cannot be settled inside the compiler. The compiler can offer the mechanism, specialised artifacts with guards, and it can offer the policy, bucket boundaries and padding rules. Which policy is right depends on the request mix, and the request mix is a property of the deployment. The scheduling of those requests against the artifacts, including how many can be in flight while a compilation proceeds, is the part of the system that has to absorb whatever the compiler gets wrong.