Reading the memory report

When a Pallas kernel asks for more VMEMThe TPU’s software-managed vector scratchpad, about 128 MiB. Blocks must be staged here before compute touches them; what is resident is what your schedule staged.taught in /l/tpu → than the chip has, the compiler doesn't guess or silently truncate anything. It stops and prints a report, and that report turns out to be one of the most literal error messages anywhere in this stack. Reading it line by line turns a wall of unfamiliar numbers into a short, specific list of which allocations to shrink and by how much, instead of a guessing game.

Scoped allocations are listed by size, largest first, and each one is scoped, meaning it only needs to stay live for part of the kernel's execution rather than the whole thing. The compiler already tries to reuse space between allocations that are never live at the same time, so the total in the budget line ends up smaller than the sum of every allocation would suggest. When that reuse still isn't enough to fit the request, the overflow report below is what tells you.

RESOURCE_EXHAUSTED: XLA:TPU compile permanent error. Ran out of memory in memory space vmem. Used 128.00M of 127.94M vmem. Exceeded vmem capacity by 64.0K.

Program vmem requirement 128.00M:
 scoped 128.00M

 Largest program allocations in vmem:

 1. Size: 64.00M
    Operator: op_name="jit(wrapped)/pallas_call"
    Shape: u8[67108864]{0}
    Unpadded size: 64.00M
    Tag: input window allocation for operator input 0. The window shape is f32[4096,4096], while the full shape is f32[4096,4096]. This allocation is single buffered.

 2. Size: 64.00M
    Operator: op_name="jit(wrapped)/pallas_call"
    Tag: output window allocation for operator output 0. The window shape is f32[4096,4096], while the full shape is f32[4096,4096].

Two labels do most of the explaining here. A line reading window allocation is not a compiler abstraction, it's a BlockSpec turned into physical bytes, the block size given to the kernel's grid now reserved as an actual region of VMEM. Two window allocations of matching size sitting next to each other usually mean the double buffer is working as intended, one block being computed on while its neighbor streams in from HBM. A line marked single buffered means the opposite happened: the compiler found room for only one copy of that operand, so the load of the next block can't overlap with compute on the current one.

The last line is the verdict, and it's worth reading exactly as written: 128.00M requested against 127.94M available. VMEMThe TPU’s software-managed vector scratchpad, about 128 MiB. Blocks must be staged here before compute touches them; what is resident is what your schedule staged.taught in /l/tpu → on this chip tops out around 128 MiB, and the kernel asked for a small amount more than that ceiling. The compiler refuses the allocation rather than letting the kernel spill silently into HBM. That's the correct failure mode, since it surfaces at compile time rather than during a run that finishes but runs slow.

None of this requires guessing at hardware limits from memory. The chip's VMEMThe TPU’s software-managed vector scratchpad, about 128 MiB. Blocks must be staged here before compute touches them; what is resident is what your schedule staged.taught in /l/tpu → ceiling is fixed, and the report states it in the same units it uses for your request, so the fix is arithmetic: shrink the block spec on one operand until the requested line drops under the available line, then rerun and check whether the single-buffered line disappears too.