MSA Overview

本页所有地址适用于来自 libtpu-0.0.40-cp314 wheel 的 libtpu.so（build-id 89edbbe81c5b328a958fe628a9f2207d）。其他版本会有所不同。

摘要

Memory-Space Assignment (MSA) 是一个编译期 pass，它为每个 HloValue 在其 live range 内决定应驻留在哪个物理内存层级，即 kDefault（HBM，容量充裕的层级）还是 kAlternate（稀缺的片上层级：TensorCore 上的 VMEM，以及存在时的 CMEM），然后插入 async copies，在正确的逻辑时间把 buffers 在层级之间移动。它是 XLA 中类似寄存器分配的机制，只是这里的“寄存器”是数十 MiB 的 VMEM，“spill” 是 HBM↔VMEM DMA，而错误决策的代价由 HBM bandwidth 支付，而不是 register pressure。MSA 是 TPU compile pipeline 的 Phase 7：它在 layout assignment 之后、scheduling 之前，运行于已经完成 layout 且 fused 的 HLO module 上，因为 scheduler 必须根据 MSA 已经选好的 memory spaces，以 cycles 计价 spill/refill。

熟悉 LLVM 的读者可以先保留两个类比，然后观察它们如何失效。MSA 是一个在 GlobalDecreasingSizeBestFitHeap<HloValue> 上的 coloring 问题：buffers 按 memory-boundedness score 排名，最大的、最热的 buffer 最先放入带有 multi-space 变化的单一 best-fit heap。这是贪心路径，也是几乎每个 module 的 hot path 上实际运行的内容。第二个类比是上游 XLA 文献描述的 "ILP variant"：一个关于（value, space, time）placement booleans 的 mixed-integer program。这个 solver 不在该二进制中。 libtpu 在 ILP 旗号下交付的是 MemoryBoundLoopOptimizer (MBLO)，一个逐 loop 的 sub-pass，它展开 software-pipelined loop 的两个相邻 steady-state iterations，并用带 even/odd coupling 的贪心 heap 解决其耦合放置。从任何 MSA symbol 都无法到达 operations_research::CpSolver、MPSolver、glop 或 CP-SAT 调用；链接进 libtpu 的 OR-Tools 服务于 SparseCore 和 autotuner，从不服务于 MsaAlgorithm。

这是 MSA 子集群的导览页。它确定该 pass 的角色，命名构建 OSS Options 并按 chip generation 门控 pass 的 jellyfish driver，以重实现者需要的层级展示 Run → Finish 控制流的形状，精确界定 greedy-vs-MBLO 分野（以及为什么 "ILP" 是 loop-local 而非 module-global），并映射子集群的其余部分。它不复现逐 buffer 的 AllocateSegment 级联、逐 generation 数值默认值，或 HBM/reservation policy；这些各自有页面，并在下方交叉引用。

对于重实现，导览契约是：

内存空间词汇和抽象 HBM/VMEM 分离。 kDefault/kAlternate 的物理含义，以及为什么 MSA 的整个工作都围绕稀缺层级。
jellyfish driver 及其门控。 RunMemorySpaceAssignment 从逐版本 TpuCompilationEnvironment 构建 OSS Options，并且只有当 IsMemorySpaceAssignmentEnabled 认为 chip generation 具有该层级时才运行 pass。
贪心核心：Run → FindAllocationSequence → MsaAlgorithm::Finish。 按 memory-boundedness 排名的 best-fit heap，先划出 cross-program prefetch 和 scoped reservations。
ILP-vs-greedy 分野。 MBLO 是在 two-iteration window 上带 coupling 的贪心，而不是经典 MILP；唯一真正的 "solver" 表面是离线 IOR autotuner，其结果作为 packed sort-order permutation 被重放。


Jellyfish driver	`xla::jellyfish::RunMemorySpaceAssignment` @ `0x12fc3080`（Phase 7 entry）
逐版本 gate	`xla::jellyfish::IsMemorySpaceAssignmentEnabled` @ `0x12fc1280`
Options builder	`xla::jellyfish::ComputeMemorySpaceAssignmentOptions` @ `0x12fc1440`（55 个 per-family overrides）
OSS pass driver	`MemorySpaceAssignment::Run` @ `0x1dc2e200`（vtable `0x21d1b5d8`）
HeapSimulator entry	`MsaAlgorithm::Finish` @ `0x1dc5b560`（vtable `0x21d1b8f0`）
Heap base class	`xla::GlobalDecreasingSizeBestFitHeap<xla::HloValue>`（multi-space best-fit）
"ILP" sub-pass	`MemoryBoundLoopOptimizer`；由 `MsaAlgorithm::IdentifyAndOptimizeMemoryBoundLoops` `0x1dc4b520` 驱动，运行 `MemoryBoundLoopOptimizer::Optimize` `0x1dcb9760`（496-byte object，`operator new(0x1f0)`）
Even/odd heap	`LoopOptimizerBestFitHeap::FindEvenAndOddAllocationBetween` @ `0x1dcb5580`
External solver	MSA hot path 上没有，没有 CpSolver/MPSolver/glop 可从任何 MSA symbol 到达
IR 层级	XLA HLO (`HloModule`/`HloInstruction`)，HLO 逻辑时间
置信度	CONFIRMED（字节锚定），除非某行或标注另有说明

MSA 决定什么

目的

MSA 为每个 HloValue live range 的每个 segment 分配 memory space，并插入实现这些分配的指令。对算法重要的两个空间是抽象的：

kDefault，容量充裕的层级，物理上是 HBM。每个 buffer 都能放在这里；把一个值放在 kDefault 永远不会因容量失败。MSA 的终端 fallback 是“留在 HBM”。
kAlternate，稀缺的片上层级，在 TensorCore 上物理为 VMEM（在有 CMEM 的 generation 上还包括 CMEM）。这是 MSA 真正配给的层级。当读取某个值的 op 受 HBM bandwidth 限制时，该值会受益于 kAlternate，因为把 operand 保持在片上可以消除 bandwidth stall。

二进制的 verifier 接受的词汇比这个二分法更宽：kHbm、kPinnedHbm、kVmem、kSmem、kCmem、kSflag、kHost，以及若干 SparseCore 和 BarnaCore 私有空间；但核心决策始终是“是否使用稀缺层级”。kPinnedHbm 与 kHbm 是不同空间，用于 runtime 必须锁定的 buffers（peer-DMA inputs）；repacker 绝不能重定位它们。kHost 是 host-RAM spill target，经由 host-offload custom calls 到达，而不是经由 heap。详细的逐空间容量和约束矩阵位于 msa-reservation-hbm-policy.md。

NOTE — "Alternate" 是算法对稀缺层级的名称；"VMEM" / "CMEM" 是物理实现。贪心 heap 从不看到 "VMEM"，它看到的是一个由 jellyfish driver 根据 chip 的 VMEM word count 定出的 [0, max_size_in_bytes) 字节范围。在最新 TensorCore generation（top jump-table entry，tpu_version == 5）上，该范围是每 TensorCore 64 MiB，512-byte word-aligned；CMEM 不存在（CmemSize == 0），因此 CMEM tier 处于 dormant 状态。见 msa-per-version-defaults.md。

两个输出

该决策为每个（value, segment）产生一个 Allocation 对象，每个对象携带一个 xla::HeapSimulator::Chunk{offset, size}（或对于纯 async pieces 不携带）：

text

PinnedAllocation        value stays resident in one space for the segment
CopyAllocation          async copy in/out  (kCopyStart … kCopyDone)
SlicedCopyAllocation    multi-slice async copy (latency split across windows)
WindowPrefetchedAllocation   tile-by-tile streaming read into VMEM during the op
ScopedAllocation / ReservedAllocation   bounded / pre-reserved alt-mem

这些 chunks 稍后会在 MemorySpaceAssignment::Process 中 lowering 为嵌入 compiled program 的 ProgramMemoryMetadata proto entries；load time 时，runtime 的 tpu::BestFitAllocator 会确定性地复原相同 offsets（见 memory/hbm-allocator.md 和 memory/on-device-compaction.md）。MSA 是一个静态规划器：它选择的每个 offset 都被冻结进 executable。

Jellyfish Driver {#the-jellyfish-driver}

目的

xla::jellyfish::RunMemorySpaceAssignment (0x12fc3080) 是 per-chip compilation environment 与 OSS XLA pass 之间的 TPU adaptor。它是 compile-phases.md 命名为 Phase 7 的 symbol。它的工作有两项：决定 MSA 是否在这个 chip generation 上运行，并在委托给 OSS driver 前，从 per-family flags 构建 OSS Options 结构体。decompiled body 确认它为 gate 调用 IsMemorySpaceAssignmentEnabled，然后调用 OSS MemorySpaceAssignment::Run；其 call graph 中没有任何 external solver。

Entry Point

text

xla::jellyfish::RunMemorySpaceAssignment            0x12fc3080
  ├─ IsMemorySpaceAssignmentEnabled(Target, env, module)  0x12fc1280   ── per-version gate
  │     reads Target+0x398 (tpu_version); jump table @ 0xae09ac8 (6 entries)
  ├─ ComputeMemorySpaceAssignmentOptions(...)            0x12fc1440   ── 55 OverwriteFieldIfNotDefault
  │     each option: per-family flag (vf/jf/gf/tpu_cmem) with xla_msa_ global fallback
  └─ MemorySpaceAssignment::Run(module, live_range,            0x1dc2e200  ── OSS pass driver
                                alias_analysis, alias_info, options)

算法

// xla::jellyfish::RunMemorySpaceAssignment  (0x12fc3080)
Status RunMemorySpaceAssignment(MsaOptions opts, Target target,
                                AliasInfo* alias_info, HloModule* module):
    // 1) Gate by chip generation.  AUTO -> enabled iff tpu_version >= 2
    //    (Pufferfish-class and later); per-family flag can override.
    if (!IsMemorySpaceAssignmentEnabled(target, env, module))      // 0x12fc1280
        return OkStatus();          // MSA simply does not run on this chip

    // 2) Materialise the OSS Options from the per-version environment.
    //    55 fields, each read from a per-family flag (xla_{vf,jf,gf}_vmem_*
    //    / xla_tpu_cmem_*) with an xla_msa_ global fallback.  This sizes
    //    max_size_in_bytes from the chip's VMEM word count, sets the overlap
    //    ratios, outstanding-copy caps, repack/retry limits, MBLO options.
    Options options = ComputeMemorySpaceAssignmentOptions(target,    // 0x12fc1440
                                                          alias_info, *module);

    // 3) Delegate to the unmodified OSS XLA pass driver.
    return MemorySpaceAssignment::Run(module, live_range,            // 0x1dc2e200
                                      alias_analysis, alias_info, options);

QUIRK — per-version gate 将内部 TpuVersion 0 和 1 都映射到 Jellyfish flag，且 AUTO 规则 "tpu_version >= 2" 会禁用这两个最低版本的 MSA。MSA 实际上从 Pufferfish-class generations 起启用。两个最新 generations（tpu_version 4 和 5）都路由到 Ghostlite gf flag namespace，0xae09ac8 处的 jump table 将两个索引发送到同一 family。假设每个 generation 恰好对应一个 flag family 的重实现会错误门控最新 chip。精确 version→flag 表见 msa-per-version-defaults.md。

Function Map

Function	Address	Role
`RunMemorySpaceAssignment`	`0x12fc3080`	Phase 7 entry；gate + options build + delegate
`IsMemorySpaceAssignmentEnabled`	`0x12fc1280`	Per-version gate；`tpu_version` at `Target+0x398`，jump table `0xae09ac8`
`ComputeMemorySpaceAssignmentOptions`	`0x12fc1440`	55 个 `OverwriteFieldIfNotDefault` per-family overrides
`OverwriteFieldIfNotDefault`	`0x1d73f360`	单个 per-family flag → option-field writer
`MemorySpaceAssignmentLambdaCollection::IsAsyncSliceImplemented`	`0x12fd64a0`	交给 OSS 的 TPU cost/aliasing/slice callbacks

贪心核心 — Run to Finish

目的

MemorySpaceAssignment::Run (0x1dc2e200) 是 OSS pass driver；它构建 heap algorithm，通过 HeapSimulator 运行，然后用选定的 allocations 重写 module。实际放置逻辑位于 MsaAlgorithm，它是 GlobalDecreasingSizeBestFitHeap<HloValue> 的子类；Finish() (0x1dc5b560) 是 HeapSimulator dispatch 到的函数。这是 MSA 的核心，并且它是贪心 best-fit heap，不是 solver。

Entry Point

text

MemorySpaceAssignment::Run(module, live_range, alias_analysis, ...)   0x1dc2e200
  ├─ FindAllocationSequence(live_range, alias_analysis)              0x1dc3b060
  │     ├─ builds the MsaAlgorithm instance
  │     ├─ HeapSimulator::Run(MsaAlgorithm-as-HeapAlgorithm, ...)
  │     │       └─ MsaAlgorithm::Finish()                            0x1dc5b560   ── all placement here
  │     └─ pulls out the chosen allocations_ vector
  └─ Process(live_range, alias_analysis)                            0x1dc2f5e0
        ├─ SimplifyGraph()              cleans tuple/GTE chains       0x1dc334c0
        ├─ SetSchedule()                inserts async copy start/done 0x1dc356a0
        ├─ ExportAndColorBuffers()      writes kAlternate/kDefault    0x1dc37820
        └─ ScheduleAsynchronousCopies() orders copy starts            0x1dc327e0

算法

Finish() 先划出特殊层级（cross-program prefetch、scoped reservations、colored buffers），然后可选运行 loop optimizer，随后按 memory-boundedness 排序顺序遍历剩余 buffer intervals，为每个（value, use）构建 AllocationRequest 并调用 AllocateSegment：

// MsaAlgorithm::Finish  (0x1dc5b560) — greedy best-fit core
Status MsaAlgorithm::Finish():
    // 1) Reserve cross-program-prefetch budget at the bottom of alt-mem,
    //    and (if xla_tpu_msa_reduce_scoped_vmem_limit) the scoped reservation.
    long alt_mem_top = ReserveAlternateMemoryForScopedMemoryAllocations();  // 0x1dc5b460

    // 2) Place already-decided block prefetches first (cross-program +
    //    preset assignments inherited from a prior compile or IOR replay).
    AllocateAndScheduleExistingBlockPrefetches(alt_mem_top);                // 0x1dc52960
    CreateNewBlockPrefetches(alt_mem_top);                                  // 0x1dc57660

    AllocateReservedScopedAllocations();                                    // 0x1dc62280
    ProcessColoredBuffers();      // frontend memory_space=kAlternate pins   // 0x1dc50d40

    // 3) Rank every HloValue by the memory-boundedness comparator.
    SortedBufferIntervals sorted =
        GetSortedBufferIntervals();   // MemoryBoundednessBufferIntervalComparator (0x1dcd3100)

    // 4) THE "ILP" SUB-PASS — only if MBLO is enabled for this chip.
    if (options_.memory_bound_loop_optimizer_options.enabled())
        IdentifyAndOptimizeMemoryBoundLoops();                              // 0x1dc4b520

    // 5) Greedy walk: each interval gets an AllocationRequest, then a
    //    per-buffer placement attempt.
    for (BufferInterval* interval : sorted_no_loops) {
        auto values = CreateAllocationValuesFromColocatedIntervals({interval});  // 0x1dc6fc00
        AllocateAllocationValues({values});                                // 0x1dc65600
        //   dispatches each use to AllocateSegment / Prefetch / WindowPrefetch
        //   AllocateSegment (0x1dc73ca0) is the 6-stage cascade — see
        //   msa-allocate-segment.md
    }
    FinalizeAllocations({values_});                                        // 0x1dc6ae80
    return OkStatus();

目标函数是隐式的，没有写入任何 solver。它存在于三个位置：sort key（0x1dcd3100 处的 MemoryBoundednessBufferIntervalComparator，按 memory_boundedness_score × size 排名，使最热且最大的 buffers 先行）、interval picker 中的逐 prefetch cost（在 HBM bandwidth 空闲时调度 copy），以及 MBLO 的 CalculateExecutionTime（对一个 loop period 的 wall-clock estimate）。disassembly 中没有保留 lambda/mu weighting constants；它们被烘焙进 CostAnalysis 的 benefit 和 bandwidth-idle 计算。

Function Map

Function	Address	Role
`MemorySpaceAssignment::Run`	`0x1dc2e200`	OSS pass driver：FindAllocationSequence + Process
`MemorySpaceAssignment::Process`	`0x1dc2f5e0`	用选定 allocations 重写 module
`MsaAlgorithm::Finish`	`0x1dc5b560`	贪心 best-fit core（HeapSimulator entry）
`MsaAlgorithm::AllocateAllocationValues`	`0x1dc65600`	逐 value use dispatch
`MsaAlgorithm::AllocateSegment`	`0x1dc73ca0`	6-stage 逐 buffer 级联，见 msa-allocate-segment.md
`MemoryBoundednessBufferIntervalComparator` ctor	`0x1dcd3100`	sort key（memory-boundedness × size）
`GlobalDecreasingSizeBestFitHeap::GetSortedBufferIntervals`	`0x1e48f400`	为贪心遍历按序排列 intervals

"ILP" Sub-Pass — MemoryBoundLoopOptimizer

目的

MemoryBoundLoopOptimizer (MBLO) 是 libtpu 中最接近 MSA "ILP variant" 的部分，也是该 pass 中最名不副实的部分。它不是 mixed-integer program。它是逐 loop 的贪心 allocator，通过展开 software-pipelined loop 的两个相邻 iterations 来建模 steady state，并在私有 even/odd best-fit heap 上求解它们耦合的 buffer placement。它由 options_.memory_bound_loop_optimizer_options.enabled() 门控，并在外层贪心遍历之前由 MsaAlgorithm::IdentifyAndOptimizeMemoryBoundLoops (0x1dc4b520) 调用。

为什么它“形似 ILP”但不是 ILP

经典 XLA-MSA ILP 会为每个（value, space, time）引入一个 Boolean placement variable，并把 constraint system 交给 CP-SAT 或 MILP solver。该二进制没有这样的 variable，也没有这样的 solver 调用。对每个 MSA decompile file 的验证显示：从 MsaAlgorithm、MemoryBoundLoopOptimizer、LoopOptimizerBestFitHeap 或 MemorySpaceAssignment::Run 无法到达 operations_research::CpSolver、CpModelBuilder、MPSolver、LinearExpr、glop::LinearProgram 或 math_opt::*。libtpu 中存在的约 80 个 OR-Tools/CP-SAT symbols 服务于 SparseCore minibatch packing、gRPC routing 和 autotuner，而不是 MSA。

MBLO 拥有的、替代 ILP 的是一条真正的耦合约束，使它带有 ILP-like flavor：even/odd chunk pair。当 prefetch 跨越 iteration boundary（copy 在前一个 period 开始，并在当前 period 完成）时，同一个 buffer 必须在 even-iteration heap 中占用一个 chunk，且在 odd-iteration heap 中占用另一个不同 chunk，并且二者都必须与 steady state 中其他 live 内容共存。LoopOptimizerBestFitHeap::FindEvenAndOddAllocationBetween (0x1dcb5580) 原子地返回 (even, odd) chunk pair，或同时回滚二者。这就是约束；其余部分是贪心。

Entry Point

text

MsaAlgorithm::IdentifyAndOptimizeMemoryBoundLoops          0x1dc4b520
  └─ per kWhile candidate (opcode 0x40):
     MsaAlgorithm::OptimizeMemoryBoundLoop(start, end, size) 0x1dc4a140
       └─ MemoryBoundLoopOptimizer::Create(...)            0x1dcb5c40   ── operator new(0x1f0) = 496 B
            ├─ Initialize()                                0x1dcb5d60   ── classify LoopValues
            └─ Optimize()                                  0x1dcb9760
                 = SortLoopValues + AllocateLoopValues + CalculateExecutionTime

算法

// MemoryBoundLoopOptimizer::Optimize  (0x1dcb9760)
void Optimize():
    SortLoopValues();                  // stable_sort by bandwidth score (0x1dcca880)
    AllocateLoopValues();              // 0x1dcb9840 — 5-case allocation_type dispatch
    float t = CalculateExecutionTime(); // 0x1dcbb9a0 — wall-clock estimate over one period

// AllocateLoopValues  (0x1dcb9840) — dispatch by LoopValue.allocation_type
void AllocateLoopValues():
    for (LoopValue& v : loop_values_):
        CHECK(v.allocation_type != kUnsupported);     // log-check; filtered earlier
        switch (v.allocation_type):                   // cmp 4; jmp [table]
            case kTemporary: AllocateTemporary(v);  break;  // 0x1dcbdb40
            case kPinned:    AllocatePinned(v);     break;  // 0x1dcbe180
            case kPrefetch:  collect for batch;     break;  // kLoopCarriedDependency folds into kPinned
    AllocatePrefetches(prefetch_subset);              // 0x1dcbe580 — score-sorted batch
        // per prefetch: AllocatePrefetch (0x1dcc0160) ->
        //   LoopOptimizerBestFitHeap::FindEvenAndOddAllocationBetween (0x1dcb5580)
        //     CreateEvenAllocationBlock (0x1dcb5280) + CreateOddAllocationBlock (0x1dcb5400)
        //     FindAndCommitChunkCandidate(even) ; FindChunkCandidate(odd)
        //     if !disjoint(even, odd): RemoveEvenChunks + RemoveOddChunks (rollback)

当 AllocatePrefetch 失败（没有 disjoint even/odd pair 能放下）时，该 buffer 会简单地从 MBLO plan 中丢弃；外层贪心遍历会在下一次排序迭代中把它作为普通 interval 接管。fallback 是隐式的，MBLO 返回 partial schedule，MsaAlgorithm 完成其余部分。

GOTCHA — MBLO 的窗口恰好是两个 loop periods（outer_start, outer_start + loop_size, outer_start + 2·loop_size），且 prefetch start index 可以为负（begin_idx_in_loop ∈ [-loop_size, +loop_size]），意思是“在前一个 period 开始”。把 loop body 建模为单个 iteration 的重实现会无法分配任何跨 boundary prefetch，而这些正是重点，因为它们是 steady state 用 compute 隐藏 HBM latency 的方式。两个 heaps（even, odd）存在的目的正是让相邻 iterations 的同一 buffer copies 永不 alias。

Function Map

Function	Address	Role
`IdentifyAndOptimizeMemoryBoundLoops`	`0x1dc4b520`	扫描 `kWhile` candidates，驱动 MBLO
`OptimizeMemoryBoundLoop`	`0x1dc4a140`	逐 loop wrapper；创建并应用 MBLO
`MemoryBoundLoopOptimizer::Create`	`0x1dcb5c40`	`operator new(0x1f0)` + ctor + Initialize
`MemoryBoundLoopOptimizer::Optimize`	`0x1dcb9760`	SortLoopValues + AllocateLoopValues + cost
`AllocateLoopValues`	`0x1dcb9840`	5-case `allocation_type` jump table
`AllocatePrefetch`	`0x1dcc0160`	每个 prefetch 的 even/odd coupling
`LoopOptimizerBestFitHeap::FindEvenAndOddAllocationBetween`	`0x1dcb5580`	原子 `(even, odd)` chunk-pair allocation
`CalculateExecutionTime`	`0x1dcbb9a0`	Wall-clock estimate（一个 period + async-copy wait）

NOTE — 精确的 LoopValue.allocation_type classifier（一个 buffer 何时是 kTemporary、kPinned 或 kPrefetch）位于 MemoryBoundLoopOptimizer::Initialize (0x1dcb5d60, ~3 kB) 内，尚未完全跟踪；只解码了消费该 classification 的 dispatch table。应将 classifier boundary 视为 LOW confidence。

唯一真正的 Solver — IOR Autotuner Replay

唯一有实际 optimal solver 进入 MSA 的地方在进程外。实验性 IOR（"Integer-Optimal-Repacker"-style autotuner，flag xla_msa_experimental_ior_algorithm）不会在 libtpu 内解决任何问题。它从磁盘加载预先计算的 solution（xla_ior_stored_solution_path）并重放。重放表面是单个 proto：

text

message MemorySpaceAssignmentConfig {   // descriptor @ 0xbed7cdc
  repeated uint64 order = 1 [packed = true];   // a permutation of buffer-interval ranks
}

因此，"stored solution" 只是一个packed permutation of buffer-interval indices，也就是 sort order。离线 solver 的输出是一个排名，被反馈给同一个贪心 best-fit heap 作为预排序顺序。进程内 MSA 始终是贪心 heap。Flags xla_ior_use_stored_solution、xla_ior_fast_mem_run_production_msa 和 xla_ior_fast_mem_round_trip_production_msa 控制加载和 replay 后验证。Flags xla_msa_experimental_use_telamalloc、xla_tpu_msa_use_minimalloc、xla_tpu_msa_use_tinymalloc 选择 ILP-style repackers，但它们没有静态链接进 libtpu；它们是为兼容 Google 内部构建而保留的 inert toggles。

MSA 子集群地图

Page	Owns	Key anchors
msa-overview.md（本页）	角色、jellyfish driver、贪心 `Finish`、ILP-vs-greedy framing	`0x12fc3080`, `0x1dc5b560`, `0x1dcb9760`
msa-allocate-segment.md	6-stage 逐 buffer 级联（required-assign → pin → no-copy → prefetch → evict → default）；`AllocationRequest`（约 0x210 B）布局；`AllocationResult` 11-bit 失败位掩码	`0x1dc73ca0`, `0x1dc72820`, `0x1dc7bc80`
msa-per-version-defaults.md	Version→flag gating matrix（`tpu_version` 4 和 5 共享 `gf`）、5-variant per-family override scheme、逐 generation overlap ratios / outstanding-copy caps、最新 generation 上 64-MiB VMEM / 512-B-word budget	`0x12fc1280`, `0x12fc1440`, table `0xae09ac8`
msa-reservation-hbm-policy.md	逐 memory-space capacity & constraint matrix、scoped + cross-program-prefetch reservations、`kPinnedHbm`/`kHost` policy、HBM-allocator handoff	`0x1dc5b460`, `0x1dc62280`

prefetch interval picker（CostAnalysisPrefetchIntervalPicker, ctor 0x1dcd6b60; Begin 0x1dcd7a00 / Next 0x1dcd7e20）、async-copy time-bucket resource model（AsynchronousCopyResource::HasEnoughResource 0x1dc78d40），以及 cost model（CostAnalysis::Create 0x1dceafc0）由 AllocateSegment 咨询，并记录在 msa-allocate-segment.md。

Component	Relationship
layout-assignment.md	Phase 6；固定 MSA 大小所依赖的 physical tile layouts；紧接在 MSA 之前运行
compile-phases.md	记录 Phase 7 placement 和完整有序 phase spine
memory/vmem-allocator.md	MSA 的 `kAlternate` chunks 落入的 runtime VMEM tier
memory/hbm-allocator.md	复原 MSA 静态 offsets 的 runtime HBM allocator
sched/overview.md	Phase 8 scheduler；根据 MSA 分配的 spaces 计价 spill/refill
sched/lhs-ilp-variant.md	一个不同的 "ILP variant"，属于 latency-hiding scheduler，而不是 MSA

交叉引用

msa-allocate-segment.md — Finish 驱动的逐 buffer placement cascade 和 request/result layout
msa-per-version-defaults.md — jellyfish driver 读取的 gating matrix 和 numeric defaults
msa-reservation-hbm-policy.md — per-space capacity、reservations 和 HBM allocator handoff
compile-phases.md — 有序 TPU compile spine 中的 Phase 7
layout-assignment.md — Phase 6，作为前驱，其 layouts 决定 MSA buffers 的大小
overview.md — Part V orientation；MSA 作为 XLA 的 HBM/VMEM register allocator
sched/overview.md — Phase 8，MSA memory-space decisions 的消费者
memory/hbm-allocator.md — 复原 MSA frozen offsets 的 runtime allocator
memory/vmem-allocator.md — 物理 kAlternate tier

MSA Overview ​

摘要 ​

MSA 决定什么 ​

目的 ​

两个输出 ​

Jellyfish Driver {#the-jellyfish-driver} ​

目的 ​

Entry Point ​

算法 ​

Function Map ​

贪心核心 — Run to Finish ​

目的 ​

Entry Point ​

算法 ​

Function Map ​

"ILP" Sub-Pass — MemoryBoundLoopOptimizer ​

目的 ​

为什么它“形似 ILP”但不是 ILP ​

Entry Point ​

算法 ​

Function Map ​

唯一真正的 Solver — IOR Autotuner Replay ​

MSA 子集群地图 ​

相关组件 ​

交叉引用 ​

MSA Overview

摘要

MSA 决定什么

目的

两个输出

Jellyfish Driver {#the-jellyfish-driver}

目的

Entry Point

算法

Function Map

贪心核心 — Run to Finish

目的

Entry Point

算法

Function Map

"ILP" Sub-Pass — MemoryBoundLoopOptimizer

目的

为什么它“形似 ILP”但不是 ILP

Entry Point

算法

Function Map

唯一真正的 Solver — IOR Autotuner Replay

MSA 子集群地图

相关组件

交叉引用