VMEM 分配器

本页中的所有地址、vtable 偏移和字段偏移均适用于 libtpu-0.0.40-cp314 wheel 中的 libtpu.so（build-id 89edbbe81c5b328a958fe628a9f2207d）。其他版本会有所不同。

摘要

VMEM（vector memory）是 TensorCore 稀缺的片上暂存区：它是 MSA 配给的层级，是 HBM 与 MXU/VPU 之间的中转区，也是编译器负担得起保留在片上的每个溢出值的归宿。本页记录 VMEM arena：如何按芯片世代划出一段平坦字节范围、每次分配都会向上取整到的 tile-alignment quantum、每代 VMEM 大小，以及在允许 MSA 触碰其余空间之前，编译器如何为 scoped scratch 和 MXU overlay buffers 预留 arena 的一片区域。

重新实现者必须先内化一个结构性事实：此二进制中没有名为 VmemAllocator 的类。 VMEM 由与所有其他层级相同的双栈机制服务。编译时，XLA 通过 MsaAlgorithm 在 GlobalDecreasingSizeBestFitHeap<HloValue> 上放置 kVmem 值；加载时，运行时把冻结的偏移重新水合到一个泛型 tpu::BestFitAllocator，它以 MemoryAllocator::Config{base_offset=0, end=VmemSizeBytes, alignment, granule} 为 VMEM 层级实例化。分配器路径中唯一特定于 VMEM 的代码，是少数按世代提供大小、对齐量子、bank 数和预留尾部的 xla::jellyfish::*Target::Vmem* 虚函数覆盖。本页下文记录的全部内容，就是这些虚函数加上消耗它们的预算算术。

本页负责的 arena 契约：

Arena 范围。 每个 TensorCore 为 [0, VmemSizeBytes)，base_offset 始终为 0，大小读取自芯片部件 proto 在启动时填充的一个 int32 字段。
对齐量子。 按世代而异：Jellyfish 上为 tile ChunkBytes，Pufferfish / Viperfish / Ghostlite 上为 max(GranuleBytes, VmemWordSizeBytes)。绝不是固定的编译期常量。
按世代的大小输入。 VmemSizeBytes、VmemWordSizeBytes、ChunkBytes、bank 数，以及默认 scoped budget，全部由活动的 Target 子类决定。
MSA 运行前切出的预留。 OverlayReservedVmemBytes（MXU overlay 尾部）和 ChunkBytes * GetReservedVmemBufferSizeChunks（collective staging），二者都从 arena 中扣除，以得到可用的 scoped-VMEM limit。

本页不覆盖：HBM 分配器及其合并规则，见 hbm-allocator.md；MSA 放置循环以及控制 VMEM 驻留的 gflag 开关，见 ../compiler/msa-overview.md 和 ../compiler/msa-per-version-defaults.md；片上层级图，见 overview.md。


Arena 范围	每个 TensorCore 为 `[base_offset=0, end=VmemSizeBytes)`
大小来源	`Target::VmemSizeBytes` @ `0x1d615e00`（int32 @ `Target +0x458`，符号扩展）
对齐（JF）	`JellyfishTarget::VmemAlignmentBoundaryInBytes` @ `0x1d490d40` -> `Target::ChunkBytes`
对齐（PF/VF/GL）	`*Target::VmemAlignmentBoundaryInBytes` -> `max(GranuleBytes, VmemWordSizeBytes)`
Tile 量子	`Target::ChunkBytes` @ `0x1d619f40` = `4 * topology.word_count`
Granule	`Target::VmemWordSizeBytes` @ `0x1d617300`（uint32 @ `Target +0x50C`）
Overlay 尾部	`Target::OverlayReservedVmemBytes`（vtable `+0x220`），base 为 `0`，GL 为 `16*ChunkSizeBytes`
Scoped budget	`scoped_memory_util::ScopedVmemLimitBytes` @ `0x1c864dc0`
编译期放置器	`MsaAlgorithm` 在 `GlobalDecreasingSizeBestFitHeap<HloValue>` 上（MS = `kVmem` = 3）
运行时分配器	泛型 `tpu::BestFitAllocator`（按层级 `Config`），与 HBM 共享
`VmemAllocator` 类	不存在 - VMEM 使用泛型分配器 + 按世代的 `Target` 虚函数
置信度	已确认（字节锚定），除非某行或标注另有说明

VMEM Arena

每个 TensorCore 一个平坦字节范围

VMEM 是每个 TensorCore 上单个连续的字节 arena。编译期放置器和运行时分配器都把它视为 [0, VmemSizeBytes)。基址偏移始终为零：每个芯片都从 sub-tile 地址 0 开始 VMEM，因此 VMEM 层级的运行时 Config.base_offset_in_bytes_ 被硬编码为 0，不同于 HBM，后者由运行时切出用户预留前缀。

Target::VmemSizeBytes（0x1d615e00）是一次字段读取，从 int32 符号扩展：

// xla::jellyfish::Target::VmemSizeBytes  @ 0x1d615e00
__int64 Target::VmemSizeBytes(Target *this) {
  return *((int *)this + 278);   // Target +0x458, signed int32
}
```text

已确认：`278 * 4 = 0x458`，且 IDA 类型为 `int`，因此该值被符号扩展到 64 位（所以负 sentinel 可表示；下方的 override flag 使用了一个）。该字段在启动时由 `TpuChipParts` / `TpuMemoryParts` 填充（从嵌入的 `chip_parts.binarypb` 解码而来），不是计算得出。

> **注意 -** 编译器可以在启动时用 `xla_tpu_override_vmem_size_kib` flag（句柄 `0x223a0980`）替换 arena 大小，该值读取一次并左移 10（KiB -> bytes）。它的 sentinel 为 `-1`/unset，在这种情况下逐字使用 `Target` 提供的大小。这是 arena 的 `end` 不同于芯片部件值的唯一方式。

### Arena 从低到高容纳什么

MSA 看到的可用 arena **小于** `VmemSizeBytes`。在确定 MSA 的 heap 大小前，会先扣除两项预留，见 [预留](#msa-之前切出的预留)。逻辑上：

```text
 VMEM byte range  [0 ........................................ VmemSizeBytes)
 ┌──────────────────────────────────┬──────────────┬──────────────────────┐
 │  MSA-placed values + scoped       │ collective   │  OverlayReservedVmem  │
 │  scratch (the rationed region)    │ staging      │  (MXU operand overlay)│
 │  ← GlobalDecreasingSizeBestFit →  │ chunks       │  ← off-limits to MSA →│
 └──────────────────────────────────┴──────────────┴──────────────────────┘
   usable = ScopedVmemLimitBytes  =  VmemSizeBytes
                                     − OverlayReservedVmemBytes
                                     − ChunkBytes * GetReservedVmemBufferSizeChunks

相比之下，运行时分配器看到的是完整的 [0, VmemSizeBytes) 范围；它只是重放 MSA 已经在可用子区域内选择好的偏移，所以永远不会撞到这些预留。

对齐量子

每次 VMEM 分配都会向上取整到按世代定义的对齐边界。这是核心的 tile 对齐规则，而且它不是编译期常量：它是 VmemAlignmentBoundaryInBytes 虚函数，通过活动 Target 子类在 vtable +0x5C8 处分派。

Jellyfish：tile chunk

JellyfishTarget::VmemAlignmentBoundaryInBytes（0x1d490d40）是到 Target::ChunkBytes 的纯 thunk：

// xla::jellyfish::JellyfishTarget::VmemAlignmentBoundaryInBytes  @ 0x1d490d40  (thunk)
__int64 JellyfishTarget::VmemAlignmentBoundaryInBytes(JellyfishTarget *this) {
  return Target::ChunkBytes(this);
}
```text

`ChunkBytes` 是 tile 量子：每个 topology word 四字节：

```c
// xla::jellyfish::Target::ChunkBytes  @ 0x1d619f40
__int64 Target::ChunkBytes(Target *this) {
  return 4LL * *(_QWORD *)(*((_QWORD *)this + 119) + 424LL);   // 4 * topology[+0x3B8].word_count
}

已确认：this[119] 是 topology 指针（+0x3B8）；+424（topology 结构内的 +0x1A8）是 sub-lane word_count；结果为 word_count * 4。这是填满一个 VPU bundle 的最小 tile：lane x sub-lane 量子按 4 字节 word 缩放。32 位形式为 Target::ChunkSizeBytes（0x1d617100），用于需要 uint32 的位置。

Pufferfish / Viperfish / Ghostlite：granule 与 word 中较大的一个

三个较新的 Target 共享同一函数体：对齐为分派得到的 granule 和 VMEM word size 中较大的一个：

// xla::jellyfish::{Pufferfish,Viperfish,Ghostlite}Target::VmemAlignmentBoundaryInBytes
//   @ 0x1d4952e0 / 0x1d49b8e0 / 0x1d4985c0   (byte-identical bodies)
__int64 PufferfishTarget::VmemAlignmentBoundaryInBytes(PufferfishTarget *this) {
  __int64 v1 = Target::GranuleBytes(this);                 // vtable[+0x5C0] dispatch
  __int64 result = (int)Target::VmemWordSizeBytes(this);   // uint32 @ Target +0x50C
  return v1 > (int)result ? v1 : result;
}
```text

Pufferfish（`0x1d4952e0`）、Viperfish（`0x1d49b8e0`）和 Ghostlite（`0x1d4985c0`）均已确认：三者反编译函数体逐字节相同。`GranuleBytes`（`0x1d617f80`）本身是按世代的虚函数（`vtable[+0x5C0]`）；`VmemWordSizeBytes`（`0x1d617300`）是在 `Target +0x50C` 的直接字段读取。

### Granule 和 word

`Target::VmemWordSizeBytes` 是每 lane 的 sub-word，即运行时 `Config` 使用的分配 granule：

```c
// xla::jellyfish::Target::VmemWordSizeBytes  @ 0x1d617300
__int64 Target::VmemWordSizeBytes(Target *this) {
  return *((unsigned int *)this + 323);   // Target +0x50C, uint32
}

已确认：323 * 4 = 0x50C。word size 和 GranuleBytes 都在启动时从 TpuChipParts 填充；按 codename 的数值仍等待 chip_parts.binarypb 解码（见注意事项）。

世代	对齐公式	Tile 量子（`ChunkBytes`）	Granule（`VmemWordSizeBytes`）
Jellyfish (v2)	`ChunkBytes`	`4 * topology.word_count`	chip-parts
Pufferfish (v4)	`max(GranuleBytes, VmemWordSizeBytes)`	`4 * topology.word_count`	chip-parts
Viperfish (v5)	`max(GranuleBytes, VmemWordSizeBytes)`	`4 * topology.word_count`	chip-parts
Ghostlite (v6e)	`max(GranuleBytes, VmemWordSizeBytes)`	`4 * topology.word_count`	chip-parts

注意 - Arena 的 granule（运行时分配器量化到的 Config.granule_in_bytes_）是 VmemWordSizeBytes，而对齐是上面的较大值公式。在较新的世代上二者可能不同：值的起始偏移按对齐边界取整，但其大小按 granule 取整。基类 Target::VmemAlignmentBoundaryInBytes（0x1d61e940）是不会返回的纯虚错误路径；没有 codename 子类的 Target 是 bug。

按世代的大小输入

Arena 完全由活动的 Target 子类参数化。VMEM 的数字字节大小不在 .text 中，而是通过 VmemSizeBytes 暴露的 chip_parts.binarypb 字段。代码中固化的是其他所有内容：bank 数、默认 scoped budget，以及 overlay 预留公式。

Bank 数

MemBanks(MemorySpace)（vtable +0xC0）返回每个层级的 bank 数；MemorySpace::kVmem == 3。函数体已由反编译确认：

// xla::jellyfish::JellyfishTarget::MemBanks  @ 0x1d48fc80
// MS==3 (kVmem) → 8 ;  MS==5 (kSmem) → 2 ;  else LogFatal
// xla::jellyfish::GhostliteTarget::MemBanks @ 0x1d4969c0
// MS==3 (kVmem) → 32 ;  MS==5 → 8 ;  else LogFatal
// xla::jellyfish::PufferfishTarget::MemBanks @ 0x1d493900
// qword_B5305C8[MS-3]  for MS ∈ {3,4,5}  → {16, 32, 8} ;  else LogFatal
```text

Jellyfish、Pufferfish、Viperfish、Ghostlite 均已确认。Pufferfish 针对连续范围 `MS ∈ {3..5}` 索引 `0xb5305c8` 处的 3 项 rodata 表（`{16, 32, 8}`）。Viperfish（`0x1d4999c0`）返回 VMEM=32、MS=5->8（反编译确认，形状与 Ghostlite 相同）。

| Target | VMEM banks (MS=3) | `kSmem` (MS=5) | Cross-slot bank conflicts |
|---|---:|---:|:---:|
| JellyfishTarget | 8 | 2 | false |
| PufferfishTarget | 16 | 8 | false |
| ViperfishTarget | 32 | 8 | **true** |
| GhostliteTarget | 32 | 8 | **true** |

Banking 是*访问调度*属性，而不是分配属性：分配器发放字节偏移，LLO bundle packer 在发射时推导 `(bank, sub-bank) = (offset / VmemWordSizeBytes) mod MemBanks(kVmem)`。cross-slot-conflict 位（Viperfish/Ghostlite 为 `true`）驱动 `xla_jf_avoid_cross_slot_vmem_bank_conflicts` swizzle 插入；它不会改变 arena 布局。

### 默认 scoped-VMEM budget

`DefaultPlatformScopedMemoryBytes`（vtable `+0x228`）是在 limit 算术钳制之前、每个程序的 scoped scratch 高水位默认值：

```c
// JellyfishTarget::DefaultPlatformScopedMemoryBytes  @ 0x1d48fc40  →  0x1000000  (16 MiB)
// GhostliteTarget::DefaultPlatformScopedMemoryBytes  @ 0x1d497540  →  0x2000000  (32 MiB)

已确认。Pufferfish（0x1d494520）和 Viperfish（0x1d49a720）均返回 0x1000000（16 MiB），已由反编译确认。基类 Target 版本（0x1d61d200）是 LogMessageFatal；每个具体世代都必须覆盖它。

Target	DefaultPlatformScopedMemoryBytes
JellyfishTarget	16 MiB (`0x1000000`)
PufferfishTarget	16 MiB
ViperfishTarget	16 MiB
GhostliteTarget	32 MiB (`0x2000000`)
`Target::` base	`LogMessageFatal` - 永不抵达

编译 flag xla_tpu_scoped_vmem_limit_kib（句柄 0x223b8770）在存在时覆盖此默认值（在 DefaultScopedVmemBytes @ 0x1c864e40 内，从偏移 0x10F0 的 TpuCompilationEnvironment proto 读取）；-1 sentinel 选择按 Target 的默认值，否则 proto 值 << 10（KiB -> bytes）。

MSA 之前切出的预留

MSA 看到 heap 之前，会从 arena 中移除两片区域。可用 scoped-VMEM limit 由 scoped_memory_util::ScopedVmemLimitBytes（0x1c864dc0）计算，反编译确认其精确为：

// xla::jellyfish::scoped_memory_util::ScopedVmemLimitBytes  @ 0x1c864dc0
__int64 ScopedVmemLimitBytes(/*Target*/ this, const Target *a2, const HloModule *a3) {
  __int64 v3 = Target::VmemSizeBytes(this);
  __int64 v4 = (*(vtable[+0x220]))(this);                 // OverlayReservedVmemBytes
  __int64 v7 = 0;
  if (a2) {
    __int64 v5 = Target::ChunkBytes(this);
    __int64 env = GetTpuCompEnv(a3);
    v7 = v5 * ring_sum_emitter_utils::GetReservedVmemBufferSizeChunks(this, env);
  }
  return v3 - (v7 + v4);
}
```text

已确认：函数体计算 `VmemSizeBytes - OverlayReservedVmemBytes - ChunkBytes * GetReservedVmemBufferSizeChunks(comp_env)`。（注意，被调用的 vtable 槽是 `+0x220`；原始 `+0x5D0` 标注指的是*声明的* `OverlayReservedVmemBytes` 虚函数，二者都指向运行时通过 Target vtable 到达的同一 hook。）

### Overlay 预留（MXU 操作数 staging 尾部）

`OverlayReservedVmemBytes` 从 VMEM 顶部切出一个固定尾部，用于 MXU 驻留的 "overlay" 操作数缓冲区（ping-pong staging 加 MSA 临时 scratch）。基类返回 `0`；Ghostlite 预留 16 个 tile chunk：

```c
// xla::jellyfish::GhostliteTarget::OverlayReservedVmemBytes  @ 0x1d497520
__int64 GhostliteTarget::OverlayReservedVmemBytes(GhostliteTarget *this) {
  return 16LL * (int)Target::ChunkSizeBytes(this);   // 16 tile chunks
}

Ghostlite 已确认。基类 Target::OverlayReservedVmemBytes（0x1d48fc20）返回 0；Jellyfish 和 Pufferfish 继承它。Viperfish（0x1d49a6c0）按 codename 设门：生产版 Viperfish 预留 16 * ChunkSizeBytes，但 viperfish-lite 禁用 overlay（返回 0），由对 variant string 前四字节的 cmpl $0x6574696c（"lite"，小端）测试选择（由反编译和反汇编确认：函数体先按 inline-vs-heap string 位分支，然后把开头 dword 与 0x6574696c 比较）。

Target	OverlayReservedVmemBytes
`Target::` base / Jellyfish / Pufferfish	`0`
ViperfishTarget	`16 * ChunkSizeBytes`（生产版）；若 codename 为 `lite` 则为 `0`
GhostliteTarget	`16 * ChunkSizeBytes`

Collective staging chunks

第二项预留是 ChunkBytes * GetReservedVmemBufferSizeChunks(comp_env)（ring_sum_emitter_utils::GetReservedVmemBufferSizeChunks @ 0x1c86a820）：这是由 comp-env 驱动的 tile chunk 数，为 ring-sum / all-reduce staging buffers 留出空间，因此无论 MSA 如何打包 arena 的其余部分，collective lowering 总能有一个受保证的 bounce buffer。

双栈分配路径

VMEM 采用与所有其他层级相同的双栈模式分配。两个栈都不包含特定于 VMEM 的分配器代码；二者都是由上面的 Target 虚函数参数化的泛型引擎。

编译时（XLA）

逐 HLO 放置通过 MSA 运行。判定“这个 HloValue 位于 VMEM（kAlternate）还是 HBM（kDefault）”的是 MsaAlgorithm，它运行在 GlobalDecreasingSizeBestFitHeap<HloValue> 上；这是 MSA 用于任何 alternate-memory 层级的同一个 heap，VMEM 只是 TensorCore 上的 kAlternate 实现。贪心 heap 从不看到 "VMEM"；它看到的是 [0, ScopedVmemLimitBytes) 字节范围。完整的放置循环、prefetch headroom 逻辑和 gflag 开关属于 ../compiler/msa-overview.md 和 ../compiler/msa-per-version-defaults.md；本页只负责这些 pass 从中取用的 arena。

Scoped（逐指令）scratch 由 LloRegionBuilder::AllocateScopedVmem（0x1d5182c0）分配，这是一个五条指令的 trampoline：

// xla::jellyfish::LloRegionBuilder::AllocateScopedVmem  @ 0x1d5182c0
LloValue *AllocateScopedVmem(LloRegionBuilder *a1, Shape *a2, ...) {
  return AllocateScopedMemory(a1, a2, 3u, ...);   // MemorySpace::kVmem == 3
}
```text

已确认：它以字面量 `kVmem`（3）memory-space tag 转发到泛型 `AllocateScopedMemory`。所选偏移被冻结到编译后程序中嵌入的 `ProgramMemoryMetadata_Allocation{memory_space=kVmem, offset, size, ...}` proto 条目。

### 加载时（runtime HAL）

`ProgramMemoryAllocator::CreateFromProto` 为每个内存层级重新水合一个 `tpu::BestFitAllocator`。VMEM 层级接收泛型 32 字节 `MemoryAllocator::Config`：

```text
MemoryAllocator::Config {
  base_offset = 0,                               // VMEM always starts at sub-tile 0
  end         = Target::VmemSizeBytes(),         // or xla_tpu_override_vmem_size_kib << 10
  alignment   = Target::VmemAlignmentBoundaryInBytes(),
  granule     = Target::VmemWordSizeBytes(),
}

分配器类、free-list、合并和 best-fit 算法都与 HBM 层级相同；没有 VMEM 子类。分配器内部见 hbm-allocator.md。因为 MSA 预先计算了每个偏移，所以运行时 VMEM 分配器几乎从不被要求搜索：它重放冻结的偏移，只在动态 scoped scratch 和少量 DMA staging buffers 上使用 free-list 路径。

注意 - 运行时 VMEM 层级没有 deferred-free 路径。tpu::DeferredTpuAllocator 只包装用户可见的 HBM buffers；VMEM deallocation 会内联合并。

耗尽

编译时 VMEM 耗尽是硬错误，而不是 spill：片上没有可溢出到的地方。三种失败模式会触碰 arena：

请求的 scoped VMEM 超过 limit。 IsRequestedScopedVmemValid（0x12fcbec0）返回一个 Status，由 CatPieces 从三个 rodata 字面量组装而成：" bytes of scoped Vmem requested (via "（0xa215531）、"), but the max valid bytes is " 和 ". See go/scoped-vmem for more details."；其中插入请求大小、出问题的 op 和 ScopedVmemLimitBytes()，也就是请求超过了上面计算出的可用 arena。
MSA 无法放置某个值。 MSA 最多重试 xla_jf_vmem_max_retries 次，然后把该值退回到 kDefault（HBM 驻留）。如果该值必须留在 alternate memory 中，则编译通过 xla::error::CompileTimeScopedVmemOom（0x1c62e5a0）硬失败。
运行时 OOM。 与 HBM 一致：BestFitAllocator::Allocate 返回带碎片转储的 absl::ResourceExhaustedError（见 hbm-allocator.md）。由于偏移是静态的，VMEM 很少走到这里。

Fusion-overflow rejection（CostModel::FusionWouldExceedVmemCapacity @ 0x130c4a80）是静默的 fusion 形成否决，而不是分配错误：它从一开始就阻止会超过 arena 的 fusion 形成，因此归在 fusion cost model 文档中，而不是本页。

注意事项和待办项

按 codename 的 VMEM 数字字节大小不在 .text 中。 VmemSizeBytes、VmemWordSizeBytes 和 GranuleBytes 都读取启动时 chip_parts.binarypb 解码填充的字段。本页中的公式是字节精确的；按 codename（Jellyfish、Pufferfish、Viperfish、viperfish-lite、Ghostlite）的字面字节数仍等待 proto 解码，本页不作断言。
"vN" 标签是 TpuVersion enum int，在 switch tables 中观察到（kJellyfish == 2，...）。最新生产部件由通过不同 variant_name() codename string 选择的 GhostliteTarget 服务；0.0.40 中没有单独的 GhostfishTarget 类，xla_gf_vmem_* flag 系列会重新配置同一个 MsaAlgorithm。
四个 Viperfish 函数体都已由反编译确认：alignment 虚函数（0x1d49b8e0）、MemBanks（0x1d4999c0）、DefaultPlatformScopedMemoryBytes（0x1d49a720）以及 overlay lite-branch（0x1d49a6c0）都直接读取自 Viperfish 反编译，并与反汇编交叉核对，而不是仅从 sibling 函数体推断。

交叉引用

overview.md - 五个片上层级 + host memory；VMEM 在层级结构中的位置。
hbm-allocator.md - VMEM 层级共享的泛型 tpu::BestFitAllocator 内部机制（合并、free-list、OOM 消息）。
smem-scalar-memory.md - scalar-memory 层级（MemorySpace::kSmem == 5）、分配器和寻址。
tpu-buffer-layout.md - 所选 VMEM 偏移描述的设备端 buffer 结构。
../compiler/msa-overview.md - 通过 kAlternate 配给 VMEM arena 的编译期放置器（MsaAlgorithm）。
../compiler/msa-per-version-defaults.md - 按世代的 MSA 数值默认值和 xla_jf_vmem_* flag 系列。
../compiler/layout-assignment.md - 在 MSA 之前运行并固定 VMEM 偏移必须遵守的 tile layouts 的 pass。

VMEM 分配器 ​

摘要 ​

VMEM Arena ​

每个 TensorCore 一个平坦字节范围 ​

对齐量子 ​

Jellyfish：tile chunk ​

Pufferfish / Viperfish / Ghostlite：granule 与 word 中较大的一个 ​

按世代的大小输入 ​

Bank 数 ​

MSA 之前切出的预留 ​

Collective staging chunks ​

双栈分配路径 ​

编译时（XLA） ​

耗尽 ​

注意事项和待办项 ​

交叉引用 ​