Vector memory
Also known as VMEM, TPU on-chip memory, vector memory
In plain English
VMEM is the fast on-chip scratch memory a TPU kernel works from, and how much of it a kernel uses decides how far ahead it can prefetch.
Technical definition
VMEM is the software-managed on-chip vector memory of a TPU TensorCore, the staging area that Pallas kernels fill from HBM and compute against, analogous to shared memory on a GPU.
Engineering details
A Pallas kernel hides HBM latency by double buffering: it fetches the next block into VMEM while computing on the current one, so both blocks must fit at once. The size of the compute block therefore determines how deep the prefetch can be. The vector unit reads VMEM in tiles whose last dimension is 128 lanes wide, so arrays with a smaller trailing dimension are padded up. Recurrent state for Gated DeltaNet layers stays in FP32 inside VMEM even when it is stored in BF16 in HBM, which halves the HBM footprint without changing the arithmetic.
Why it matters
VMEM pressure shows up as regressions in unexpected places. Adding a second buffer set for asynchronous GDN state transfers initially broke data-parallel attention until scratch buffers were reused. Collective offload to SparseCore is gated on whether the message fits in VMEM. Kernel authors on TPU budget VMEM the way GPU kernel authors budget shared memory and registers.
How to read it in InferenceX
The ragged paged attention v3 fix split the KV compute block from the fetch block, keeping the fetch at 16k tokens while reducing compute to 4k, which freed VMEM for the prefetch buffer and lifted decode throughput from 64.9k to 96.3k tokens per second on Qwen3-0.6B. The asynchronous GDN state transfer gained 11.3% at concurrency 512 once VMEM capacity was recovered.