Vx can model peak memory pressure at compile time

Aditya Kumar · 21 September 2026

TL;DR: Summing every tensor allocated during a Llama-2-7B prefill gives 58.93 GiB. An NVIDIA A100-40 GPU has only 40 GiB of memory. That sum makes the model look too large to fit, prompting engineers to offload layers to host RAM. In reality, the tensors resident in memory at the same time peak at 15.94 GiB, which uses 39.8% of the card. The 42.99 GiB difference is memory the program never holds simultaneously. Vx computes this peak at compile time on a laptop using the program's block structure and a declared machine model.

Deciding whether a model fits on a GPU determines your entire deployment setup: whether you need one GPU or two, resident weights or streamed weights, and your serving batch size. Engineers usually make this decision by summing all allocations. That sum measures the wrong quantity.

Why the sum is wrong

Memory capacity depends on peak residency during execution. A tensor occupies GPU memory only between its allocation and its release. Two tensors need space simultaneously only when their lifetimes overlap.

During transformer inference, memory usage splits cleanly into two groups:

• Weights stay resident. The model reads every layer to generate each token. As a result, all 32 layers of weight matrices remain in GPU memory from the first token to the last. They are all live at the same time.
• Activations are temporary. Activations are intermediate values produced during layer computation. When layer 8 begins, the attention scores from layer 7 are already freed. A naive sum adds up activations for all 32 layers, but the GPU only holds one layer of activations at any moment.

The 42.99 GiB difference comes entirely from temporary activations. Counting released activations leads to unnecessary layer offloading, which severely degrades performance.

A sum, a peak, and twenty lines

The Vx compiler tracks memory lifetimes across lexical blocks to distinguish peak usage from total allocations. Here is a minimal program demonstrating this behavior. It places one persistent weights buffer and one loop-scoped scratch buffer into a 4 MiB memory space:

Memory CPU_DRAM {}
Memory HBM {
  within: Memory::CPU_DRAM, capacity: 4 MiB, bandwidth: 1555 GB/s, managed: cached
}
Topology Dev {
  arch: nvptx64,
  memory: Memory::HBM,
  visible: [Memory::HBM],
  transfer Memory::CPU_DRAM -> Memory::HBM : 31.5 GB/s
}

fn main() -> i32 {
  let weights_h = Tensor<f32, [512, 512]>::uninit();
  let mut weights = transfer(weights_h, Memory::HBM);
  for layer in 0..32 {
    let scratch_h = Tensor<f32, [512, 512]>::uninit();
    let scratch = transfer(scratch_h, Memory::HBM);
    spawn on(Topology::Dev) {
      weights[0][0] = weights[0][0] + scratch[0][0];
    }
  }
  print(2);
  return 0;
}

The program allocates thirty-three 1 MiB buffers across the run inside a 4 MiB high-bandwidth memory (HBM) space. The compiler accepts the program and reports its memory usage:

"resident_sets": [
  {"space": "HBM", "total_bytes": 2097152, "capacity_bytes": 4194304,
   "utilization": 0.5, "tiles": 2, "overcommit": false}
]

The report shows two tiles occupying 2 MiB, which is 50% utilization. At any moment, only the persistent weights buffer and the scratch buffer for the current iteration exist simultaneously.

If you unroll the loop by placing all 32 scratch buffers at function scope, the program's lifetime structure changes. Because function-scoped variables stay allocated until main returns, all 33 buffers remain live at the same time. The compiler rejects the unrolled program:

Error[E6010]: the working set placed in memory space 'HBM' (33 tiles) sums to 34603008 bytes,
              over its 4194304 byte capacity; place fewer/smaller tiles or declare it `overcommit`

The unrolled working set requires 33 MiB against a 4 MiB capacity, reaching 8.25 utilization. Both compiler verdicts are correct. The unrolled code holds 33 buffers at once because no block boundary releases them earlier.

How the compiler knows

Vx calculates peak memory usage by tracking scope nesting and allocation order at compile time. The compiler assigns a unique identifier to each lexical scope (code block enclosed in braces) as it enters that block. During type checking, it maintains a scope chain: an ordered list of active block identifiers from outermost scope to innermost scope.

Each allocation records three attributes: its size in bytes, its active scope chain, and its sequential position in program order. At the end of a function, the compiler computes the peak memory required for each space:

peak ← 0
for each tile T, in program order:
    live ← Σ { U.bytes : U placed no later than T,
                         and U.scope is a prefix of T.scope }
    peak ← max(peak, live)

The prefix test checks for coexistence. If tile U belongs to a block that encloses T, then U is still open when T is placed. Tiles in sibling blocks were already released when those blocks closed. The program order check ensures that tiles placed in an outer scope after an inner block closes are not counted as coexisting with that inner block.

Three language design choices make this analysis computable at compile time:

• Placement is part of the type. Calling transfer(t, Memory::HBM) produces a value whose type identifies the memory space and budget it draws from.
• Spaces declare capacities. Memory spaces specify explicit capacities in the machine model, providing fixed limits instead of runtime hardware queries.
• Release follows block structure. Because the runtime reclaims memory at block exits, the compiler determines buffer overlap directly from the syntax tree.

The compiler accounts for memory at block exits. A borrow checker can track when a variable is last read, but the Vx runtime frees the backing buffer when the surrounding lexical block closes. If compile-time accounting assumed last-use reclamation, the compiler would approve programs that exceed memory limits and crash at runtime. The static check accurately models the runtime release discipline.

Llama-2-7B, for real

Applying this peak analysis to production workloads shows that Llama-2-7B fits comfortably on a single 40 GiB GPU. Consider standard Llama-2-7B dimensions: 32 layers, hidden dimension 4096, feed-forward network (FFN) intermediate dimension 11008, 32 attention heads, and a vocabulary size of 32000. The workload is a 4096-token prefill at batch size 1 in 16-bit floating point (f16), with attention scores materialized. The target machine is fleet/a100-40.vx, declaring 40 GiB of HBM at 1555 GB/s and a host link at 31.5 GB/s.

The weights are placed at function scope because they persist throughout execution. The per-layer activations are placed inside the layer loop because each iteration reclaims them:

fn main() -> i32 {
  // Resident for the whole run: embeddings, the KV cache, and one tile per
  // layer standing for that layer's seven weight matrices.
  let embed = transfer(embed_h, Memory::HBM);        // 250 MiB
  let lm_head = transfer(lm_head_h, Memory::HBM);    // 250 MiB
  let mut kv = transfer(kv_h, Memory::HBM);          // 2 GiB
  let w0 = transfer(w0_h, Memory::HBM);              // 386 MiB, x32
  ...

  for layer in 0..32 {
    // Dead at the end of the iteration, every one of them.
    let scores = transfer(scores_h, Memory::HBM);     // 1 GiB
    let qkv = transfer(qkv_h, Memory::HBM);           // 96 MiB
    let ffn = transfer(ffn_h, Memory::HBM);           // 172 MiB
    ...
  }
  return 0;
}

Testing both allocation layouts against the same machine model produces these compile-time results:

Total allocations sum to 63,271,075,840 bytes (58.93 GiB), whereas peak residency requires only 17,112,760,320 bytes (15.94 GiB). The peak breaks down into 12.06 GiB of layer weights, 0.49 GiB of embeddings, 2 GiB of key-value (KV) cache, and 1.39 GiB for a single layer's activations. The remaining 42.99 GiB represents activations from the other 31 layers that never coexist in memory.

The analysis preserves the full size of persistent data. All 32 layers of weights are live concurrently, so they are fully counted at 12.06 GiB. Memory savings come entirely from temporary activation buffers whose lifetimes do not overlap.

What the wrong answer costs

Assuming that total allocations determine fit causes an unnecessary 47× throughput penalty. If an engineer believes the model requires 58.93 GiB, they will conclude it cannot fit on a 40 GiB GPU. To run anyway, the system must stream weights from host CPU memory across the PCIe bus for each layer. The machine model prices this transfer directly because the host link is a declared edge. The compiler reports the planned transfer cost:

{"path": ["CPU_DRAM", "HBM"], "bytes": 404750336,
 "derived_cost": 12849217016, "derived_unit": "ps", "cost_source": "link_rate"}

Moving one layer's weights across the host link takes 12.85 ms. Across 32 layers, streaming requires 411 ms per token, capping throughput at 2.4 tokens per second. Keeping those same weights resident in GPU HBM allows reading all 13.48 GB at 1555 GB/s. That takes only 8.7 ms, providing a roofline throughput of 115 tokens per second.

This calculation gap creates a 47× throughput penalty. Accurate peak accounting determines whether the GPU serves the model at full memory bandwidth or stalls waiting for bus transfers.

Where the line actually is

The compile-time check enforces physical memory limits by rejecting configurations that exceed capacity. Keeping the prefill length at 4096 tokens and raising the batch size on the same A100-40 illustrates this boundary:

Batch 8 fits with 0.9% of GPU memory to spare (39.65 GiB used). Batch 9 requires 43.03 GiB and exceeds card capacity, so the compiler rejects it with error E6010. The compiler CLI (vxc) runs on a laptop and answers deployment capacity questions before hardware is provisioned or jobs are scheduled.

What it does not do

The peak analysis is intentionally conservative. The compiler guarantees that it will never approve a program that exceeds capacity. In edge cases, it may reject a program that would have fit at runtime.

• Analysis stops at the function boundary. The compiler checks memory budgets per function. A caller's active buffers are not yet counted against a callee's allocations. Calculating peak usage across the full call tree requires call graph analysis and recursion handling, tracked in Vx#444.
• Loop bodies are checked once. The compiler analyzes loop bodies a single time because the runtime reclaims iteration buffers on each pass. This check assumes loops do not accumulate buffers into outer collections across iterations.
• Only verified scopes release memory early. The compiler currently frees allocations at the end of if branches and loop bodies, matching its emitted intermediate representation (IR). Other constructs, such as spawn regions, match arms, comptime blocks, and unsafe blocks, keep their allocations alive until the enclosing function returns.
• Shapes must be static. If a tensor dimension is not known at compile time, the compiler issues a warning and omits that allocation from the compile-time budget check.

Why this is not the usual state of affairs

Mainstream toolchains cannot compute peak working sets at compile time because their type systems lack hardware memory bounds.

In C, C++, or Rust, heap allocations are not tied to bounded memory spaces. Because compilers in those languages do not track capacity limits, programs discover memory exhaustion only when an allocator returns null. In CUDA, static shared memory (__shared__) is scoped to individual kernels and does not compose across lexical scopes, and device global memory has no compile-time budget. Machine learning frameworks discover tensor shapes and device targets at runtime dispatch, leaving peak memory detection to the runtime allocator.

Runtime discovery remains the standard approach today, but runtime estimation errors are costly. An overestimate leads to unnecessary host offloading or larger GPU rentals that add latency to every token. An underestimate causes out-of-memory crashes in production. By combining typed placement, declared hardware capacities, and lexical lifetimes, Vx calculates exact peak requirements before provisioning hardware.

The analysis implementation is described in docs/working_set_peak_implementation_plan.md and tested in tests/integration_test/working_set_peak_test.rs. Test cases include tests/frontend/pass/working_set_is_a_peak_not_a_sum.vx and tests/frontend/fail/three_dtcm_tiles_do_not_fit_a_cortex_m7.vx. Cross-function peak analysis is tracked in Vx#444. Vx is licensed under Apache 2.0 with the LLVM exception. You can install Vx, or read What a pointer forgets to learn why memory spaces belong in the type system.