VMEM（向量存储）
也称为 VMEM、vector memory、TPU 片上存储
先用大白话
VMEM 是 TPU kernel 直接读写的片上高速暂存区，kernel 占用多少 VMEM 决定了它能提前预取多远。
技术定义
VMEM 是 TPU TensorCore 上由软件管理的片上向量存储，Pallas kernel 从 HBM 往这里加载数据并在此计算，作用类似 GPU 的 shared memory。
工程细节
Pallas kernel 通过双缓冲隐藏 HBM 延迟：在当前块计算时把下一块预取进 VMEM，所以两个块必须同时放得下，计算块的大小决定了预取深度。向量单元按最后一维为 128 lane 的 tile 读取 VMEM，末维更小的数组会被补齐。Gated DeltaNet 层的循环状态在 HBM 中以 BF16 存储，但在 VMEM 内仍以 FP32 计算，这使 HBM 占用减半而算术精度不变。
为什么重要
VMEM 压力会在意想不到的地方表现为性能回退。为异步 GDN 状态传输增加第二组缓冲区时，一度导致数据并行注意力回退，直到复用了 scratch buffer 才恢复。集合通信是否 offload 到 SparseCore 也以消息能否装进 VMEM 为门槛。TPU kernel 作者预算 VMEM 的方式，与 GPU kernel 作者预算 shared memory 和寄存器一样。
如何在 InferenceX 中解读
ragged paged attention v3 的修复把 KV 计算块与预取块分开，预取块保持 16k token，计算块缩到 4k，为预取缓冲释放出 VMEM，使 Qwen3-0.6B 的 decode 吞吐量从每秒 64.9k token 提高到 96.3k。VMEM 容量恢复后，异步 GDN 状态传输在并发 512 下带来 11.3% 的收益。