MSA 预留与 HBM 策略

本页中的所有地址、符号、偏移、proto 名称和魔数常量均适用于 libtpu-0.0.40-cp314 wheel 中的 libtpu.so（build-id 89edbbe81c5b328a958fe628a9f2207d，构建 libtpu_lts_20260413_b_RC00）。其他版本会有所不同。

摘要

TpuCustomCallMemorySpacePolicy 是一个 HLO pass，它按 custom-call 决定缓冲区落在哪里：要么为 memory-space-assignment（MSA）优化器预留一片 VMEM，要么强制 custom-call 的缓冲区进入 HBM。它消耗的旋钮是单个 proto 子消息 TpuCustomCallMemorySpaceSpec，作为 TPU 编译环境上的 AutoOr<spec> 字段携带，因此用户要么固定一个 spec，要么保持 AUTO 并让编译器合成一个。

本页规范化反编译所钉住的两个字节精确事项：

两个策略字段字典。 TpuCustomCallMemorySpaceSpec 是 oneof policy { MsaReservationPolicy msa_reservation_policy = 1; HbmPolicy hbm_policy = 2; }。MsaReservationPolicy 只携带一个字段：uint64 msa_reservation_size_bytes = 1；HbmPolicy 是零字段标记（选择它本身就是指令）。两个默认实例都是全零字节，并由生成的 _table_、_globals_ 和序列化的 DescriptorProto 三重来源确认。
ResolveMemorySpaceSpec @ 0x11036320。 该解析器总是先合成一个候选 MsaReservationPolicy，其大小由目标 VMEM 计算，然后仅当 spec 为 AUTO/缺失时保留该候选，否则逐字返回用户的 spec。AUTO 大小为 max(floor(free_vmem / 4), 10 MiB)，其中 free_vmem = VmemSizeBytes − (OverlayReservedVmemBytes + DefaultScopedVmemBytes)。

下游分派（RunImpl @ 0x110364a0）按解析后的 oneof case 切换：case 1 运行由 msa_reservation_size_bytes 定大小的 MSA 预留；case 2 强制 HBM 放置；case 0 为 no-op。整个 pass 由 memory-space assignment 是否启用来门控（ShouldEnablePass @ 0x110362a0）。


Pass 类	`xla::jellyfish::TpuCustomCallMemorySpacePolicy`
旋钮 proto	`TpuCustomCallMemorySpaceSpec`（oneof `policy`），arm 1 `MsaReservationPolicy`，arm 2 `HbmPolicy`
旋钮载体	编译环境上的 `AutoOr<TpuCustomCallMemorySpaceSpec>`（comp-env `+0xbb0` = `+2992` 处的 `AutoProto*` 槽）
解析器	`ResolveMemorySpaceSpec` @ `0x11036320`（0x16f B）
分派器	`RunImpl` @ `0x110364a0` — 按解析后的 oneof case 切换
门控	`ShouldEnablePass` @ `0x110362a0` → `IsMemorySpaceAssignmentEnabled` @ `0x12fc1280`
MSA arm 动作	`RunMsaReservationPolicy` @ `0x110367c0`
HBM arm 动作	`RunHbmPolicy` @ `0x11038120`
AUTO 大小公式	`max(floor((Vmem − Overlay − ScopedVmem) × 0.25), 10 MiB)`
魔数常量	`0.25f`（`.rodata` @ `0x84a2724` = `0x3e800000`）；10 MiB 下限（`0x00a00000`）
源文件	`platforms/xla/service/jellyfish/tpu_custom_call_memory_space_policy.cc`（从日志字符串恢复）
置信度	HIGH（符号/字符串/字节锚定），除非某行或标注另有说明

策略旋钮：`TpuCustomCallMemorySpaceSpec`

TpuCustomCallMemorySpaceSpec 是按 custom-call 放置旋钮的结构化 arm。它是一个名为 policy 的单一 oneof，带有两个 arm；选择一个 arm 同时选择了放置动作并提供其（唯一）参数。wire 结构通过三种独立方式恢复，且字节一致：

.data.rel.ro 中生成的 <Msg>::_table_ TcParseTableBase + FieldEntry 行，
.data 中的 <Msg>_globals_ 默认实例字节，以及
嵌入 protodesc_cold 的序列化 google.protobuf.DescriptorProto（tpu_compilation_environment.proto 的 message_type 子记录）。

protobuf

// xla.jellyfish (recovered, byte-exact)
message TpuCustomCallMemorySpaceSpec {
  oneof policy {                                       // oneof_decl[0] = "policy"
    MsaReservationPolicy msa_reservation_policy = 1;   // oneof_index 0
    HbmPolicy            hbm_policy             = 2;    // oneof_index 0
  }

  message MsaReservationPolicy {
    uint64 msa_reservation_size_bytes = 1;             // the ONLY field of either arm
  }

  message HbmPolicy {}                                 // zero fields — a pure marker
}
```text

两个字段共享 `oneof_index 0`，因此它们**互斥**：一个 spec 要么是 MSA 预留，要么是 HBM 放置，绝不会两者兼有；默认两者都不是（`case 0`，unset）。因此这个结构化旋钮收缩为三种状态之一：*unset*（让 AUTO 合成）、*为 MSA 预留 N 字节 VMEM*，或*强制 HBM*。

### 字段字典（`_table_` ⇔ `_globals_` ⇔ `DescriptorProto`）

`c++off` 是字段存储在消息对象内的字节偏移；`has_idx` 是 hasbit 索引；`type_card` 是解析表 type-card（singular-uint64 = `0x08d1`；oneof message arms = `0x0436`）。三个消息的 `_globals_` 都是全零字节，因此 proto 默认是 *unset oneof / size 0 / HBM-marker absent*。

| 消息 | field# | 名称 | proto 类型 | c++off | has_idx | type_card | 默认 |
|---|---|---|---|---|---|---|---|
| `TpuCustomCallMemorySpaceSpec` | — | *(oneof `policy`)* | — | — | — | — | unset (case 0) |
| &nbsp;&nbsp;arm 1 | 1 | `msa_reservation_policy` | message | `0x10` (union) | case @ `0x1c` | aux0 / `0x0436` | absent |
| &nbsp;&nbsp;arm 2 | 2 | `hbm_policy` | message | `0x10` (union) | case @ `0x1c` | aux1 / `0x0436` | absent |
| `MsaReservationPolicy` | 1 | `msa_reservation_size_bytes` | uint64 | `0x18` | 128 | `0x08d1` | 0 |
| `HbmPolicy` | — | *(no fields — marker)* | — | — | — | — | — |

```text
_table_   VAs:  spec 0x21cfa708 · MSA 0x21cfa658 · HBM 0x21cfa6b8
_globals_ VAs:  spec 0x223c8920 · MSA 0x223c8538 · HBM 0x223c8558   (all byte-zero)
ctors:          spec 0x1db25fa0 · MSA via Arena::DefaultConstruct 0x1db638e0 · HBM 0x1db25ea0

[NOTE] 父 _table_ @ 0x21cfa708 携带一个 oneof：两个 arm 都存储在对象 +0x10（union），活动 case 判别器在 +0x1c。_table_+0x68 处的 aux 子表指针重定位到 MsaReservationPolicy 表（0x21cfa658）和 HbmPolicy 表（0x21cfa6b8），并由 readelf -r relocation 遍历确认。

MsaReservationPolicy::ByteSizeLong 的反编译（@ 0x1db25e60）直接佐证了这个单字段字典：它测试对象 +0x10 的 bit 0 hasbit，读取 *((qword*)this + 3)（= +0x18）处的 uint64，并加上标准单 uint64 varint 长度 (9·bsr64(v|1) + 137) >> 6。没有其他字段被序列化，因为并不存在其他字段。

解析器：`ResolveMemorySpaceSpec` @ `0x11036320`

ResolveMemorySpaceSpec(out, Target&, HloModule&, AutoOr<spec>) 是一个 0x16f 字节的 sret 函数。它的定义性行为是：它总是先构建一个按 VMEM 定大小的候选 MsaReservationPolicy，即使用户已经提供了 spec，然后才决定保留该候选还是返回用户 spec。反编译结构如下：

// 0x11036320 — recovered control flow (cleaned)
TpuCustomCallMemorySpaceSpec *ResolveMemorySpaceSpec(
    TpuCustomCallMemorySpaceSpec *out, Target *target,
    const HloModule *module, const AutoOr<spec> *spec)
{
  TpuCustomCallMemorySpaceSpec local;                 // -0x50, case = 0
  TpuCustomCallMemorySpaceSpec::ctor(&local, /*arena=*/0);

  MsaReservationPolicy *msa;
  if (local._case == 1) {                             // (never on the fresh local)
    msa = local._union_ptr;
  } else {
    local.clear_policy();
    local._case = 1;                                  // select msa_reservation_policy
    msa = Arena::DefaultConstruct<MsaReservationPolicy>(arena_of(local));
    local._union_ptr = msa;
  }

  // --- AUTO size: one quarter of the FREE vmem, floored at 10 MiB ---
  int64_t vmem    = Target::VmemSizeBytes(target);              // Target+0x458, i32 -> i64
  int64_t overlay = (*target->vptr[0x220/8])(target);          // OverlayReservedVmemBytes()
  int64_t scoped  = scoped_memory_util::DefaultScopedVmemBytes(target, module);
  int64_t free    = vmem - (overlay + scoped);
  int64_t quarter = (int64_t) truncf((float)free * 0.25f);     // mulss 0x84a2724 ; cvttss2si
  uint64_t size   = (quarter >= 0xa00001) ? quarter : 0xa00000; // max(quarter, 10 MiB)

  msa->msa_reservation_size_bytes = size;             // store at msa+0x18
  msa->_hasbits[0] |= 1;                              // hasbit at msa+0x10 bit 0

  // --- keep the AUTO candidate, or return the user's spec ---
  if (spec->present == 0) {                           // AUTO  (byte at AutoOr+0x28 == 0)
    out = TpuCustomCallMemorySpaceSpec::ctor(out, 0);
    /* CopyFrom / InternalSwap the local AUTO candidate into out */
  } else {                                            // user-supplied spec
    const spec *user = (spec->variant_idx == 1)       // variant index at AutoOr+0x20
                         ? *(spec**)spec               //   1 = pointer
                         : spec;                       //   0 = inline; else throw
    out = TpuCustomCallMemorySpaceSpec::ctor(out, 0, user);  // copy the USER spec
  }
  local.~TpuCustomCallMemorySpaceSpec();
  return out;
}
```text

[NOTE] `AutoOr<T>` 包装器在 `+0x28` 存储 presence 标志（0 = `AUTO`/absent），在 `+0x20` 存储 variant 判别器（0 = inline value，1 = pointer-to-value；其他值触发 `std::__throw_bad_variant_access`）。`RunImpl` 在调用解析器之前，通过 `AutoOr<spec>::FromProtoOrDie`（@ `0x1103f2e0`）从 comp-env 物化 `AutoOr`。

实际含义是：**AUTO 并不是“什么都不做”**，AUTO 表示“合成一个占空闲 VMEM 四分之一的 MSA 预留”。用户 spec（无论哪个 arm）会整体替换该候选。

---

## AUTO 大小公式

空闲 VMEM 的推导，带字节精确常量：

```text
free  = Target::VmemSizeBytes()                                  ; movslq [Target+0x458]   (i32 -> i64)
        − ( Target::OverlayReservedVmemBytes()                   ; (*vptr)[0x220]
            + scoped_memory_util::DefaultScopedVmemBytes(T, M) )  ; @0x1c864e40
size  = (uint64) truncf( (float)free * 0.25f )                   ; mulss [0x84a2724] = 0x3e800000 = 0.25f
size  = max(size, 0x00a00000)                                    ; cmp $0xa00001 / cmovge → 10 MiB floor

项	来源	VA / 常量
total VMEM	`Target::VmemSizeBytes`（读取 `Target+0x458`，i32）	`0x1d615e00`
− overlay reserve	`Target::OverlayReservedVmemBytes`（vtable 调用 `*0x220`）	base `0x1d48fc20` 返回 0；Ghostlite = `ChunkSizeBytes() << 4`；Viperfish = `ChunkSizeBytes() << 4`，但 `"lite"`（v5e）SKU 返回 0
− default scoped VMEM	`scoped_memory_util::DefaultScopedVmemBytes`	`0x1c864e40`
× `0.25f`	`.rodata` float	`0x84a2724`（`0x3e800000`）
trunc → uint64	`cvttss2si`	—
floor `0x00a00000`	10 MiB 下限 clamp	`cmp $0xa00001 ; cmovge`
→ `msa_reservation_size_bytes`	存储到 `MsaPolicy+0x18`，hasbit `MsaPolicy+0x10` bit 0	—

scoped-VMEM 工作集项

DefaultScopedVmemBytes（@ 0x1c864e40）是在取四分之一之前从 VMEM 中划出的工作集。没有 HloModule 时，它尾调用每目标 cap；否则它预留 ring-sum buffer 区域并 clamp 到 cap。该 cap 并不只是平台默认 vtable 调用：它是 TpuCompEnv+0x10f0 处的 comp-env 字段，以 KiB 解释（<< 10），并且只有当该字段读为 -1 时才使用 (*vptr)[0x228] 平台默认值作为 fallback：

text

capKiB  = *(int64*)(TpuCompEnv + 0x10f0)                ; signed; -1 means "unset"
cap     = (capKiB == -1) ? (*vptr)[0x228]()             ; DefaultPlatformScopedMemoryBytes() fallback
                         : (capKiB << 10)                ; else comp-env override, KiB -> bytes
total   = Target::VmemSizeBytes()
overlay = (*vptr)[0x220]()                              ; OverlayReservedVmemBytes()
chunkB  = Target::ChunkBytes()                          ; (*(Target+0x3b8))[0x1a8] * 4
chunks  = ring_sum_emitter_utils::GetReservedVmemBufferSizeChunks(Target, TpuCompEnv)   ; @0x1c86a820
scoped  = total − (chunks * chunkB + overlay)
return    (scoped >= cap) ? cap : scoped                ; min(scoped, cap)
```text

[NOTE] comp-env scoped-cap override 字段位于 `TpuCompEnv+0x10f0`，并以 KiB 为单位按有符号 `int64` 读取；把它反向映射到其命名 TCE 旋钮仍为 **OPEN**（LOW）。当该字段为 `-1` 时，每 SKU 的 `DefaultPlatformScopedMemoryBytes`（`*0x228` vtable 槽）提供 cap。

`GetReservedVmemBufferSizeChunks`（@ `0x1c86a820`）在 `inference_short_ring_sum_emitter_utils::IsSupported` 成立时返回 `16 · *(TpuCompEnv+0x12c0)`（即 `(TpuCompEnv+0x12c0) << 4`），否则返回 0，也就是说 ring-sum buffer chunk 数是 `+0x12c0` 处的 comp-env 字段。把该字段反向映射到其命名 TCE 旋钮仍为 **OPEN**（LOW）。

### 每目标常量

vtable 槽通过 **`vptr = vtable_symbol + 0x10`** ABI 偏移解析（由 `Target` 基类 ctor 在 `lea` vtable 符号后执行 `add $0x10,%rax` 证明）。朴素的“symbol+offset 处 reloc”读法会误命名这些槽；修正后的映射为：

| 调用偏移 | 解析为 | 方法 |
|---|---|---|
| `*0x220` | symbol `+0x230` | `Target::OverlayReservedVmemBytes() const` |
| `*0x228` | symbol `+0x238` | `Target::DefaultPlatformScopedMemoryBytes() const` |

`DefaultPlatformScopedMemoryBytes` 按具体 SKU 的值（每个都是 6 字节 `mov $imm,%eax ; ret`）：

| Target | scoped cap | overlay reserve |
|---|---|---|
| `JellyfishTarget` | `0x1000000` = 16 MiB | 0 |
| `ViperfishTarget` | `0x1000000` = 16 MiB | `"lite"`（v5e）SKU 上为 `0`，否则为 `ChunkSizeBytes() << 4`（v5p） |
| `PufferfishTarget` | `0x1000000` = 16 MiB | 0 |
| `GhostliteTarget` | `0x2000000` = 32 MiB | `ChunkSizeBytes() << 4` |
| `Target` (base) | string-keyed lookup @ `0x1d61d200`（LOW：未转写） | 0 |

因此在 Jellyfish/Pufferfish 上，AUTO 预留是 `max(free/4, 10 MiB)`，其中 `free = VMEM − (Overlay + ScopedVmem)`，overlay reserve 为零，scoped working set cap 为 16 MiB；Ghostlite 将 working set cap 设为 32 MiB，并加入 `16 · ChunkSizeBytes()` 的 overlay reserve。参见 [MSA Per-Version Defaults](msa-per-version-defaults.md)，了解*消耗*该预留的每代 MSA 调优。

---

## 分派器：`RunImpl` @ `0x110364a0`

`RunImpl` 是 pass 入口。它读取 comp-env，解析 spec，运行 alias analysis，然后按解析后的 oneof case 切换。恢复出的流程如下：

```c
// 0x110364a0 — recovered (cleaned)
StatusOr<bool> RunImpl(HloModule *module, const flat_hash_set<string_view> &execution_threads)
{
  Target *target = module->target();                          // [module+8]
  AutoProto *env  = GetTpuCompEnv(module)->slot(2992) ?: &AutoProto_globals_;

  AutoOr<spec> au; AutoOr<spec>::FromProtoOrDie(&au, env);     // 0x1103f2e0
  spec resolved;
  ResolveMemorySpaceSpec(&resolved, target, module, &au);     // 0x11036320
  // ... copy `resolved` into module-side spec_ (CopyFrom / InternalSwap) ...
  destroy(au);

  VLOG(1) << "Memory space spec: " << module->spec_;          // .cc:511

  switch (module->spec_._case /* [module+52] = +0x34, mirrors spec+0x1c */) {
    case 0:  return true;                                     // unset → no placement change
    default: {
      auto aa = HloAliasAnalysis::Run(module, alias_info);    // .cc:518 on error
      switch (module->spec_._case) {
        case 2: return RunHbmPolicy(this, module, aa);        // 0x11038120  (HBM)
        case 1: return RunMsaReservationPolicy(this, module, aa); // 0x110367c0 (MSA)
        default: return true;
      }
    }
  }
}

[NOTE] 解析后的 oneof case 从 [module + 52]（= +0x34）读取，它镜像 spec 的 _case_ 字段 spec+0x1c。Case 0 在 alias analysis 构建前就短路：不运行 HloAliasAnalysis::Run，立即返回 StatusOr<bool> OK。Cases 1 和 2 先构建 alias analysis，然后分派。每个 arm 还会通过 DCHECK 再次守卫自己的 oneof case（如下），因此分派不变量被检查两次。

Arm 1 — `RunMsaReservationPolicy` @ `0x110367c0`

重新守卫 DCHECK(spec_.has_msa_reservation_policy())（.cc:534），记录 "Running MSA reservation policy"，并针对 alias analysis 遍历 HloModule::MakeNonfusionComputations（@ 0x1e5e2320）。它从 spec union（MsaReservationPolicy+0x18）读取解析后的预留字节数，并划出这么多字节的 VMEM 作为 MSA 优化器可使用的预留。逐指令工作通过内部 MsaReservationPolicyInstructionProcessor 运行（指令按 IO 大小排序）。

Arm 2 — `RunHbmPolicy` @ `0x11038120`

重新守卫 DCHECK(spec_.has_hbm_policy())（.cc:578），记录 "Running HBM policy"，并且像 MSA arm 一样，针对 alias analysis 遍历 non-fusion computations。由于 HbmPolicy 没有字段，没有大小可读；该 arm 通过 MemorySpaceColorMap 无条件强制匹配的 custom-call 缓冲区进入 HBM。选择 hbm_policy arm 就是完整指令。

[NOTE] OPEN — 放置写入点与匹配谓词。 两个 arm 都已确认会遍历 MakeNonfusionComputations + alias analysis，且（MSA arm）消耗 +0x18；具体通过 color map / shape mutator 写入的 MemorySpace enum 整数，以及 custom-call 的匹配谓词（哪些 target 名称合格），没有字节追踪到 store immediate。对放置效果的置信度：MSA 预留 VMEM、HBM 强制 HBM 为 HIGH；写入的具体 enum 值为 LOW。

门控：`ShouldEnablePass` @ `0x110362a0`

该 pass 仅在模块启用 memory-space assignment 时运行。ShouldEnablePass 调用 IsMemorySpaceAssignmentEnabled(Target, ObjectView<TpuCompilationEnvironment>, HloModule)（@ 0x12fc1280）并返回其 StatusOr<bool>：

// 0x110362a0 — recovered (cleaned)
StatusOr<bool> ShouldEnablePass(/* Target, env, module */)
{
  StatusOr<bool> enabled = IsMemorySpaceAssignmentEnabled(/* ... */);
  if (enabled.ok()) {
    return {.ok = true, .value = enabled.value};   // movb $0,+8 ; movq $1,+0  (StatusOr ok rep)
  } else {
    return Status::AddSourceLocation(enabled.status(), /*line*/475,
                                     "platforms/.../tpu_custom_call_memory_space_policy.cc");
  }
}
```text

因此按 custom-call 的放置取决于更宽泛的 MSA enable flag；如果模块不运行 memory-space assignment，这个 pass 会被整个跳过，AUTO 预留也绝不会被合成。`IsMemorySpaceAssignmentEnabled` 读取的确切 TCE flag 未转写（LOW）；门控关系本身为 HIGH。

---

## 重新实现契约

忠实的重新实现必须：

1. 将 `TpuCustomCallMemorySpaceSpec` 建模为 `oneof policy`，其 arms 为 `MsaReservationPolicy { uint64 msa_reservation_size_bytes = 1; }` 和 `HbmPolicy {}`，默认 unset（case 0）。
2. 对 `AUTO`/缺失 spec，合成一个 `MsaReservationPolicy`，设置 `msa_reservation_size_bytes = max(floor((Vmem − Overlay − ScopedVmem) × 0.25), 10 MiB)`，使用单精度 float 乘法 + 截断（不是 double，也不是整数除法）以及精确的 `0x00a00000` 下限。
3. 将 `ScopedVmem` 推导为 `min(Vmem − (ring_sum_chunks · ChunkBytes + Overlay), PlatformScopedCap)`，使用每 SKU cap（JF/VF/PF 为 16 MiB，GL 为 32 MiB）和每 SKU overlay reserve。
4. 对用户提供的 spec，逐字返回它；**不要**把 AUTO 候选合并进去。
5. 按解析后的 oneof case 分派：1 → 按字段大小预留 VMEM；2 → 强制 HBM；0 → no-op。只为 cases 1/2 构建 alias analysis。
6. 用 memory-space-assignment-enabled 门控整个 pass。

---

## 交叉引用

- [MSA 概览](msa-overview.md) — 消耗本页所划预留的 memory-space-assignment ILP 优化器。
- [MSA AllocateSegment](msa-allocate-segment.md) — 消耗 VMEM budget 的分配主体 + 配置 proto。
- [MSA 每版本默认值](msa-per-version-defaults.md) — 叠加在预留之上的每代 overlap ratios / outstanding-copy caps。
- [Custom-Call Lowering 与 Target Registry](custom-call-lowering.md) — custom-call 分派层；此放置 pass 是 custom-call 元数据的一个消费者。
- [VMEM 分配器](../memory/vmem-allocator.md) — AUTO 预留据以定大小的运行时 VMEM 空间。
- [HBM 分配器](../memory/hbm-allocator.md) — `hbm_policy` arm 强制缓冲区进入的 HBM 空间。
- [内存子系统概览](../memory/overview.md) — VMEM/HBM/CMEM/SMEM 空间图。
- [返回索引](../index.md)

MSA 预留与 HBM 策略 ​

摘要 ​

策略旋钮：TpuCustomCallMemorySpaceSpec ​

解析器：ResolveMemorySpaceSpec @ 0x11036320 ​

scoped-VMEM 工作集项 ​

Arm 1 — RunMsaReservationPolicy @ 0x110367c0 ​

Arm 2 — RunHbmPolicy @ 0x11038120 ​

门控：ShouldEnablePass @ 0x110362a0 ​