HBM BestFit 分配器

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

摘要

tpu::BestFitAllocator 是唯一一个具体的 tpu::MemoryAllocator 子类，服务于 TPU core 上的每一个运行时内存层级：HBM、VMEM、SMEM、CMEM、SFLAG。它运行在一个固定的 [base_offset, allocatable_range_end) 字节跨度上，该跨度在启动时由 MemoryAllocator::Config 三元组建立。本页记录运行时 / 驱动侧 HBM 分配器：双 free-list 数据结构、best-fit 搜索、释放时积极的双向合并、split / fragmentation 策略，以及对齐量子。它不是编译器侧 placement pass；后者是 XLA 的 MSA，会把静态 offset 固化进程序，并只把这些 offset 交给本分配器重放。本页中的 free-list 算术只运行于动态表面（用户输入/输出 buffer、传输 staging、async-copy scratch、动态 program stack）；MSA 固定的静态 offset 只会被重放，从不重新推导。

没有一个类字面上叫 HbmAllocator。HBM 层级是 BestFitAllocator 的一个实例；同一个类逐字节也支撑 VMEM 和其他片上层级。每个层级（以及每个芯片代际）的唯一区别是 Config。这个分配器在结构上值得关注，因为它不像 tsl::BFCAllocator 那样使用 bin bucket（libtpu 只在主机侧使用后者，见 embedded-tcmalloc.md）。它是经典的边界标签方案：一个按 offset 做 key 的 SwissTable 保存每个 block（free 和 allocated），用于 O(1) 邻居查找；其上叠加一棵按 size 排序的 free block 红黑树，用于精确 best-fit 搜索。

熟悉 dlmalloc 的读者会认出边界标签和积极合并；这里不同的是没有 size bin，free tree 给出真正的 best-fit（最小的可容纳 block）而不是分离列表近似，分配放在所选 free block 的顶部，并且 region 永不增长，容量在构造时固定。

对重新实现而言，约定如下：

双结构。 blocks_by_offset_（Abseil flat_hash_map<offset, HashTableEntry>）是包含所有 block 的边界标签 map；free_tree_（std::set<FreeBlock, BlockOrder>）只保存 free block，并按 (size, then begin) 升序排列。
Best-fit 搜索。 FindAllocatableBlock 在 size-ordered tree 上执行 lower_bound({end − aligned, end})，返回的第一个 node 就是 size ≥ request 的最小 block，并会跳过 reserved-bottom。每实例的 Policy 开关可选择 O(n) first-fit walk。
合并。 每次 Deallocate 都立即、积极、双向合并。下一个邻居 = map.find(self.end)；上一个邻居 = 用 self.begin 做 key 的 free-tree find_equal。MergeBlock 扩展低地址 entry 并删除高地址 entry。
Split / fragmentation 策略。 分配自顶向下放置；SplitBlock 没有 minimum-remainder threshold，任何非零 leftover 都成为 free block，精确匹配则不产生 leftover。内部碎片（round-up padding）归 allocation 所有，而不是单独记账。外部碎片在 OOM diagnostic 中报告，并由 compaction 回收。
对齐量子。 请求在请求边界处一次性 round up 到 alignment_in_bytes_；region end 在构造时向下取整，因此每个 offset、size 和 remainder 都是 alignment 的倍数。HBM 的量子是 1024 B DMA 最小值（见 hbm-dma-alignment.md）。


类	`tpu::BestFitAllocator` : `tpu::MemoryAllocator`（200 字节实例，`operator new(0xC8)`）
构造函数	`0x1e817500`（`Config const&`、`Policy`）；typeinfo `0x21d346e8`，vtable `0x21d34630`
Allocate	`0x1e817820`（vt+0x30）→ `StatusOr<int64_t>`（返回 `base + offset`）
Deallocate	`0x1e819dc0`（vt+0x38）
Best-fit 搜索	`FindAllocatableBlock` @ `0x1e818540`
Split	`SplitBlock` @ `0x1e819060`；无 min-remainder
Coalesce	`MergeBlock` @ `0x1e819700`；比较器实现 `__lower_upper_bound` `0x1e824960` / `__find_equal` `0x1e824640`
Defrag	`Compact` @ `0x1e81c360` — 归 on-device-compaction.md 所有
Free list	size-ordered RB-tree `std::set<FreeBlock, BlockOrder>` + offset-keyed SwissTable `blocks_by_offset_`
HBM alignment quantum	1024 B（`jf_driver::kHbmMinimumDmaAlignment`）— 通过 `Config` 按层级设置；见 hbm-dma-alignment.md
源码标签	`learning/45eac/tpu/runtime/hal/internal/best_fit_allocator.cc`（来自 `LogMessageFatal` / error site-string）
置信度	CONFIRMED（以字节为锚点），除非某行或 callout 另有说明

范围与边界

本页是运行时/驱动侧 HBM 分配器。三个密切相关的问题各有自己的页面；不要在这里重复：

关注点	负责页面
编译期 HBM/VMEM placement（`kDefault`/`kAlternate` 决策，冻结 offset）	../compiler/msa-overview.md
1024 B HBM DMA-issue alignment check 和 16 KiB 编译期 program alignment	hbm-dma-alignment.md
OOM 触发的 `Compact` defrag plan、`TpuMemmoves`、DMA-program codegen、OOM retry chain	on-device-compaction.md
五层内存映射和主机侧分配器	overview.md、embedded-tcmalloc.md

本页负责的是：free-list 结构、best-fit 搜索、合并、split / fragmentation 策略和对齐量子，也就是一个 BestFitAllocator 实例的原地 free-list 数学。

注意 — 同一个 BestFitAllocator 类支撑 VMEM、SMEM、CMEM 和 SFLAG。算法在各层级和各芯片代际（Dragonfish 到 6acc60406）之间相同，任何分配器方法内部都没有 switch (TpuVersion) 分支。按层级和按代际的差异完全承载在 Config 三元组中，启动时来自嵌入的 chip_parts.binarypb 资源。兄弟层级见 vmem-allocator.md 和 cmem-pool.md。

两个耦合的数据结构

分配器维护两个结构，而不是一个 free list。这是核心设计事实。

blocks_by_offset_ — Abseil flat_hash_map<int64_t, HashTableEntry, OffsetHash>，位于 this+0x08。这是一个以 block offset 为 key 的边界标签 map，保存每个 block（free 和 allocated）的 {offset, state, size, metadata}。这使邻居查找成为 O(1)：物理上位于 self 上方的 block 正好是 map.find(self.end)（其 offset 等于 self 的 end）。
free_tree_ — std::set<FreeBlock, BlockOrder> 红黑树，根在 this+0x28，只保存 free block，按 (size, then begin-offset) 升序排序。在这棵树上做 lower_bound 会得到最小的可容纳 block，即带最低地址 tie-break 的精确 best-fit。

因此它是一个用于合并的地址索引边界标签方案，叠加一个用于 best-fit 搜索的 size-ordered tree。这不是 tsl::BFCAllocator 的 size-bucketed bin 方案（后者仅用于主机）。

// 200-byte instance (operator new(0xC8)); ctor 0x1e817500 sets these fields. Offsets byte-confirmed.
struct tpu::MemoryAllocator::Config {            // 32 B, passed by const&
    int64_t base_offset_in_bytes_;     // +0   >= 0  (validated; else fatal)
    int64_t allocatable_range_end_;    // +8   > 0   (capacity = end - base)
    int64_t alignment_in_bytes_;       // +16  > 0, power of two, divides granule
    int64_t granule_in_bytes_;         // +24  hardware granule (page/word)
};

enum class Policy     : int32_t { kBestFit = 0, kFirstFit = 1 };
enum class BlockState : int32_t { kFree = 0, kAllocated = 1, kReserved = 2 };

struct tpu::BestFitAllocator {                    // sizeof = 0xC8 (200)
    void*    vptr;                     // +0x00  -> vtable+0x10
    flat_hash_map<int64_t,HashTableEntry,OffsetHash>
             blocks_by_offset_;        // +0x08  boundary-tag map (ALL blocks)
    set<FreeBlock,BlockOrder>
             free_tree_;               // +0x28  RB-tree (FREE blocks only)
                                       //   size@+0x38, end_node@+0x48
    Policy   policy_a_;                // +0x40  copy 1
    /* free_tree_end_ */               // +0x48
    Policy   policy_b_;                // +0x50  copy 2 (read by Find/Split/Dealloc)
    int64_t  base_offset_in_bytes_;    // +0x58
    int64_t  alignment_in_bytes_;      // +0x60
    int64_t  max_aligned_size_;        // +0x68  INT64_MAX - (INT64_MAX % align)
    int64_t  allocatable_range_end_;   // +0x70  end - (end % align)
    int64_t  granule_in_bytes_;        // +0x78
    int64_t  bytes_in_use_;            // +0x80  } packed XMM pair, updated
    int64_t  num_allocations_;         // +0x88  } with vpaddq / vpsubq
    int64_t  peak_bytes_in_use_;       // +0x90
    int64_t  peak_block_bytes_in_use_; // +0x98
    int64_t  bytes_allocatable_init_;  // +0xA0  = end (init mirror)
    int64_t  reserved_bottom_limit_;   // +0xA8  ReserveBottomOfMemory current limit
    int64_t  peak_reserved_bottom_;    // +0xB0  max() watermark of reserved-bottom limit
    int64_t  capacity_in_bytes_;       // +0xB8  = end
    int64_t  reservation_count_;       // +0xC0
};

struct FreeBlock { int64_t lo_offset; int64_t hi_offset; };   // 16 B RB-tree value
// BlockOrder(a,b): a < b  <=>  (a.size < b.size) ||
//                              (a.size == b.size && a.lo < b.lo),
//                 where size := hi_offset - lo_offset.   (size-major, offset-minor)

struct HashTableEntry {                // 32 B SwissTable value
    int64_t offset;     // +0  = map key
    int32_t state;      // +8  kFree / kAllocated / kReserved
    int32_t _pad;       // +0xC
    int64_t size;       // +0x10  == aligned size (padding owned by allocation)
    int64_t metadata;   // +0x18  allocation id / owner
};

struct Block {                         // 48 B value materialised by GetBlockIf
    int64_t offset;     // +0
    int64_t end;        // +8   = offset + size
    int32_t state;      // +0x10
    int64_t metadata;   // +0x18
    void*   ctrl_ptr;   // +0x20  (Abseil ctrl byte, or nullptr if not found)
    void*   slot_ptr;   // +0x28  (HashTableEntry*)
};
```text

**初始状态**（ctor）：一个覆盖整个 region 的 `HashTableEntry {offset = 0, state = kFree, size = end}` 和一个 `FreeBlock {0, end}`。ctor 用一个 `operator new(0x30)` node 和一次 `__find_equal`-then-insert 给 tree 播种，然后验证 `Config`（见 [Alignment](#the-alignment-quantum)）。

### 为什么是两个结构而不是一个

单独的 offset-ordered tree 可以给出 O(log n) 邻居查找，但 best-fit 搜索是*线性的*（必须扫描以寻找最小的可容纳空洞）。单独的 size-ordered tree 可以给出 O(log n) best-fit，但*无法*找到 freed block 的物理邻居以进行合并。边界标签 map + size tree 的配对同时给出 **O(log n) best-fit 和 O(1) 摊销邻居查找**。代价是在每次 split 和 merge 时保持两者同步，这正是 `SplitBlock`/`MergeBlock`/`ShrinkFreeListEntry`/`AddToFreeList` 所做的事。

| 字段 / 结构 | 地址或偏移 | 作用 |
|---|---|---|
| `blocks_by_offset_` | `this+0x08` | 所有 block 的边界标签 map |
| `free_tree_` | `this+0x28`（size `+0x38`，end `+0x48`） | size-ordered free RB-tree |
| `policy_b_` | `this+0x50` | 由 `Find`/`Split`/`Dealloc` 读取，用于选择 best-fit vs first-fit |
| `reserved_bottom_limit_` | `this+0xA8` | 底部 watermark；搜索会跳过 |
| `BlockOrder` `lower_bound` | `0x1e824960` | 比较器 `(size, then lo)` 升序 |
| `BlockOrder` `find_equal` | `0x1e824640` | 同一比较器；prev/next 合并查找 |

---

## Best-Fit 搜索

`FindAllocatableBlock`（`0x1e818540`）是 placement 搜索。反编译确认了两个由 `policy_b_`（`this+0x50`，读作 `*((_DWORD*)this + 20)`）控制的分支：

```c
// FindAllocatableBlock(int64_t aligned) const   (0x1e818540)
Iterator FindAllocatableBlock(int64_t aligned) const {
    if (policy_ != kBestFit) {                       // FIRST-FIT (policy@+0x50 != 0)
        for (node = free_tree_.begin(); node != end(); ++node)  // in-order RB walk
            if (node.hi - node.lo >= aligned)        // node[5]-node[4] >= aligned
                return node;                         // first block that fits
        return end();
    }
    // BEST-FIT: lower_bound on the size-ordered tree.
    int64_t end = allocatable_range_end_;            // this+0x70  (read as this+14)
    FreeBlock key{ end - aligned, end };             // {lo = end-aligned, hi = end}
    auto it = free_tree_.lower_bound(key);           // __lower_upper_bound_unique_impl
    if (it == free_tree_.end()) return end();        // no fit -> caller reports OOM

    // RESERVED-BOTTOM SKIP: if >= 2 free blocks and this candidate starts exactly
    // at the bottom watermark, step past it and rescan forward for the next fit.
    if (free_tree_size_ >= 2 && it->lo == reserved_bottom_limit_) {  // +0x38, +0xA8
        for (++it; it != end(); ++it)
            if (it->hi - it->lo >= aligned) return it;
        // none after the reserved block fits -> fall through to the candidate
    }
    return it;                                       // smallest fitting block
}

关键洞见在 lower_bound key 中。tree 按 (size, then lo) 排序。用 {end − aligned, end} 作为 key，会让 key 的派生 size 正好等于 end − (end − aligned) = aligned。lower_bound 返回第一个 (size, lo) ≥ (aligned, end − aligned) 的 node，也就是第一个 size ≥ aligned 的 node。由于 tree 以 size 为主序，这个第一个 node 就是大小 ≥ 请求的最小 block。这是真正的 best-fit，而不是 bin 近似。字节 trace 确认了 v13[0] = v7 - a2; v13[1] = v7，其中 v7 = this[+0x70]，a2 = aligned。

first-fit 分支按offset顺序遍历 tree 吗？不是。它通过 std __tree iterator（result[1] left-child / result[2] parent successor 逻辑）按中序（size-then-offset）遍历，并返回第一个足够大的 block。由于 tree 按 size 排序，它到达的第一个可容纳 node 实际上也是最小的可容纳 node。因此在这个分配器上，first-fit 和 best-fit 会收敛到同一个 block，真正差异只是 first-fit 总是从最小 free block 重新扫描（最坏 O(n)），而 best-fit 通过 lower_bound 直接跳到目标（O(log n)）。所有五个 AllocatorFactoryByName factory lambda 都用 kBestFit 构造，因此本构建中设备上的 first-fit 路径处于休眠状态。

reserved-bottom watermark

ReserveBottomOfMemory（0x1e81b0c0）设置 reserved_bottom_limit_（+0xA8）。搜索会把 lo 等于该 watermark 的 free block 当成不可用，在存在其他 free block 时跳到下一个 candidate。这是运行时在该层级底部划出受保护区域的方式（XLA flag xla_tpu_user_reserved_hbm_bytes 在启动时流到这里）。free_tree_size_ >= 2 guard 确保如果 reserved block 是唯一的 free block，它仍会被返回，而不是让分配直接失败。

函数	地址	作用
`FindAllocatableBlock`	`0x1e818540`	best-fit `lower_bound` + reserved-bottom skip；first-fit walk
`__lower_upper_bound_unique_impl<FreeBlock>`	`0x1e824960`	size-ordered tree 上的 `lower_bound`
`ReserveBottomOfMemory`	`0x1e81b0c0`	设置 `+0xA8` watermark
`GetBlockIf`	`0x1e818f20`	SwissTable H2 probe → 48-B `Block`（或 null-slot Block）
`BytesAllocatable`（vt+0x70）/ `BytesAvailable`（vt+0x68）	`0x1e819b80` / `0x1e81dd60`	最大空闲连续段 / 总空闲；提供 OOM diagnostic 的 $l a r g e s t$ / $f r e e$ 值

Allocation 与 Split / Fragmentation 策略

Allocate（0x1e817820）验证 size、向上取整、寻找 block、在该 block 的顶部 split、缩小 free entry、更新 stats，并返回绝对地址。

// Allocate(int64_t size) -> StatusOr<int64_t>   (0x1e817820), byte-confirmed
StatusOr<int64_t> Allocate(int64_t size) {
    if (size < 0 || size > max_aligned_size_)              // +0x68
        return BadSizeError(size);                         // 0x1e818240

    // round up to alignment (power of two). The (size!=0) term keeps 0 -> 0.
    int64_t aligned = (size + alignment_ - (size != 0)) & -alignment_;   // +0x60

    auto it = FindAllocatableBlock(aligned);               // 0x1e818540
    if (it == free_tree_.end())                            // == this+0x48
        // message streamed (not a Substitute template), pieces verbatim:
        //   "Attempting to allocate " << HR(size)
        //   << ". That was not possible. There are " << HR(BytesAvailable())  // vt+0x68
        //   << " free. The largest contiguous region of free memory is "
        //   << HR(BytesAllocatable())                                         // vt+0x70
        //   << " due to fragmentation.";
        // ResourceExhaustedErrorBuilder, line 129; if FLAGS_tpu_log_allocations_on_oom,
        // also dumps the map.
        return ResourceExhaustedError(...);

    Block found = GetBlockIf(it->lo, kFree);               // materialise the free block
    int64_t split_offset = found.end - aligned;            // PLACE AT TOP of the block
    SplitBlock(found, split_offset, /*left=*/kFree, /*right=*/kAllocated);  // 0x1e819060
    ShrinkFreeListEntry(it, found);                        // delete or re-key the free node

    bytes_in_use_ += aligned;  num_allocations_ += 1;      // packed XMM @ +0x80
    peak_block_bytes_in_use_ = max(peak_block_bytes_in_use_, aligned);
    peak_bytes_in_use_       = max(peak_bytes_in_use_, bytes_in_use_);
    return base_offset_in_bytes_ + split_offset;           // = found.end - aligned + base
}
```text

字节 trace 精确确认了 split 调用：`SplitBlock(&out, this, &found, found.end - aligned, /*left*/0, /*right*/1)`，以及返回值 `base_offset_in_bytes_ + split_offset`。

### 自顶向下放置

分配从所选 free block 的**高端**切出：allocated range 是 `[found.end − aligned, found.end]`，而**低端** remainder `[found.begin, found.end − aligned]` 保持 free。把分配压到高端会让 region 底部在 reserved-bottom watermark 之下保持连续，从而提高小型动态分配在不扰动受保护底部区域的情况下被满足的概率。

### 没有 minimum-split-remainder

`SplitBlock`（`0x1e819060`）只有一个 guard：如果 `split_offset == orig.end`（精确匹配），它只写 left entry 并返回，不产生 remainder block。否则它总是发出 right-side entry，**无论 leftover 多小**。反编译显示 `if (a4 == orig.end) ...`（其中 `a4` 是 split offset）是唯一分支；没有与 minimum-remainder 常量比较。

```c
// SplitBlock(const Block& orig, int64_t split_off, BlockState left, BlockState right)
//   (0x1e819060) — re-states the existing slot in place, optionally adds a right entry
void SplitBlock(const Block& orig, int64_t split_off, BlockState left, BlockState right) {
    // LEFT  = [orig.begin, split_off)   gets `left`,  keeps orig.metadata
    // RIGHT = [split_off,  orig.end)    gets `right`
    if (orig.slot_ptr) {                               // re-state existing slot in place
        orig.slot_ptr->state = left;
        orig.slot_ptr->size  = split_off - orig.begin;
    } else {
        blocks_by_offset_.try_emplace(orig.begin,
            HashTableEntry{orig.begin, left, split_off - orig.begin, orig.metadata});
    }
    if (split_off == orig.end) return;                 // EXACT FIT -> no right block
    blocks_by_offset_.try_emplace(split_off,           // ANY leftover -> new entry
        HashTableEntry{split_off, right, orig.end - split_off, /*md*/0});
}

这对碎片分析很重要。min-remainder threshold（类似 dlmalloc 的 MINSIZE）会用少量内部碎片换取避免产生不可用的小 free block。这个分配器选择相反策略：它从不拒绝 split，因此从不膨胀 allocation，但可能在 tree 中留下任意小的 free block。这里可以接受，因为 remainder 永远是 alignment 的倍数（见下文），所以最小可能 free block 是一个 alignment quantum，HBM 下为 1024 B，而不是亚量子碎片。

陷阱 — 因为没有 min-remainder，长期运行且以略有不同的 size 分配/释放 buffer 的 workload 可能在 region 各处分散积累许多一量子的 free block。这是外部碎片，OOM diagnostic 中的 “The largest contiguous region of free memory is <N> due to fragmentation.” 信息正是报告它，compaction 也正是回收它。free tree 从不浪费这些空间，字节是可用的，但如果没有 defrag，就无法满足大于最大空洞的单个分配。

内部与外部碎片

内部碎片是 round-up padding：每个 allocation 的 aligned − size 字节。它归该 allocation 所有，HashTableEntry.size 保存的是对齐后 size，没有单独的 padding ledger。释放时会归还完整 aligned size。HBM 使用 1024 B 量子时，每个 buffer 最多 1023 B。
外部碎片是被拆散在非相邻 block 中的 free space。分配器用两种方式对抗它：(1) 每次 free 时积极合并（下一节）会立即重新合并相邻空洞；(2) 当即使合并后的 free space 也无法满足请求时，OOM 路径会触发 Compact（见 on-device-compaction.md）。GetFragmentation（0x1e819c80）报告 (total_free − largest_free_run) / total_free，即不位于单个最大连续段中的 free space 比例。

ShrinkFreeListEntry（0x1e819520）在 split 后保持 free tree 一致：如果 free block 被完全消耗（remaining == 0），它删除 node 并对其执行 operator delete；否则它删除并按新的排序位置重新插入 node（block 变小，因此 size key 改变）。

函数	地址	作用
`Allocate`	`0x1e817820`	round-up → find → split-at-top → shrink → stats → `base+offset`
`SplitBlock`	`0x1e819060`	left/right state；exact-fit → 无 right entry；无 min-remainder
`ShrinkFreeListEntry`	`0x1e819520`	empty 时删除 / size 变化时 remove+reinsert
`BadSizeError`	`0x1e818240`	size `< 0` 或 `> max_aligned_size_`
`GetFragmentation`	`0x1e819c80`	`(free − largest_run) / free`

Free 时合并

Deallocate（0x1e819dc0）验证地址、递减 stats、与两个物理邻居合并，并重新插入合并后的 free block。合并是积极、立即且双向的，从不在分配器层面延迟。（PJRT 层的 DeferredTpuAllocator wrapper 延迟的是对 Deallocate 的调用，不是其中的合并。）

// Deallocate(int64_t address) -> Status   (0x1e819dc0), byte-confirmed
Status Deallocate(int64_t address) {
    int64_t offset = address - base_offset_in_bytes_;        // +0x58
    HashTableEntry* slot = blocks_by_offset_.find(offset);   // SwissTable H2 probe
    if (!slot || slot->state != kAllocated)                  // state != 1 -> error
        return FailedPreconditionError(                      // line 505
            "Attempted to deallocate address %d, which is not the address of "
            "an allocated buffer.", address);

    Block self = { offset, offset + slot->size, kAllocated, slot->metadata };
    bytes_in_use_ -= slot->size;  num_allocations_ -= 1;     // packed XMM (vpsubq) @ +0x80

    // (1) COALESCE WITH PREVIOUS free block (its END == self.begin).
    //     The prev block is found by free-tree find_equal keyed on self.begin; if
    //     present, its tree node is removed and MergeBlock extends DOWNWARD.
    Block prev = GetBlockIf(self.begin - prev_size, kFree);
    if (prev.slot_ptr) {                                     // prev exists and is free
        free_tree_.erase(find_equal(FreeBlock{prev.begin, ...}));
        MergeBlock(prev, self, /*new=*/kFree, &self);        // self.begin <- prev.begin
    }
    // (2) COALESCE WITH NEXT free block (its OFFSET == self.end).
    Block next = GetBlockIf(self.end, kFree);                // map.find(self.end)
    if (next.slot_ptr) {                                     // next exists and is free
        free_tree_.erase(find_equal(FreeBlock{next.begin, ...}));
        MergeBlock(self, next, kFree, &self);                // self.end <- next.end
    }
    // (3) (re)insert the coalesced free block into the size-ordered tree.
    AddToFreeList(self);                                     // find_equal-then-insert

    VLOG(2) << "Deallocated a block start at: " << address   // line 607
            << " of " << HR(self.size) << ". Available: " << HR(BytesAvailable());
    return OkStatus();
}
```text

### 邻接检测按 offset 进行，O(1)

`blocks_by_offset_` 保存 allocated *和* free block 的全部原因就在这里：物理上位于 `self` 上方的 block，其 **offset 等于 `self` 的 end**，所以 `map.find(self.end)` 就是下一个邻居。只有当其 state 为 `kFree` 时才合并。位于 `self` 下方的 block，是 **end 等于 `self` begin** 的 free block；free tree 的 `find_equal`（用能匹配以 `self.begin` 结束的 free node 的 key）找到它，并且按构造它必然是 free（tree 只保存 free block）。两个合并都通过 `MergeBlock`。

```c
// MergeBlock(const Block& a /*lower*/, const Block& b /*upper*/, BlockState new_state, Block* out)
//   (0x1e819700) — extend the lower entry to span both; erase the upper entry.
void MergeBlock(const Block& a, const Block& b, BlockState new_state, Block* out) {
    HashTableEntry* lo = blocks_by_offset_.find(a.begin);    // keep the lower entry
    lo->size  = (a.end - a.begin) + (b.end - b.begin);       // span both
    lo->state = new_state;
    blocks_by_offset_.erase(b.begin);                        // drop the upper entry
    *out = Block{ a.begin, b.end, new_state, a.metadata, lo->ctrl, lo };
}

最多两次 merge 后，self 就是最大化合并后的 free run，AddToFreeList（0x1e81a700）会为它插入一个 FreeBlock node。净效果是：free tree 永远不包含两个相邻的 free block，邻接会当场解决。这是保持 best-fit 诚实的 invariant（否则一对碎片化相邻块的“best fit”会是错的），也是搜索中的 lo == reserved_bottom_limit_ skip 足够成立的原因。

注意 — 字节 trace 显示 prev-merge、next-merge 和 free-node（重新）插入被融合成一条代码路径，复用合并后的 self Block，并在可能时原地重设一个已有 tree node 的 key（而不是总是 erase 后再 operator new）。上面的伪代码是语义等价形式；融合形式在某个邻居已经有可扩展 tree node 时避免了一次 node 分配。

函数	地址	作用
`Deallocate`	`0x1e819dc0`	validate → dec stats → coalesce prev+next → reinsert free
`MergeBlock`	`0x1e819700`	扩展低地址 entry，删除高地址 entry
`AddToFreeList`	`0x1e81a700`	`find_equal`-then-insert `FreeBlock` RB node
`__find_equal<FreeBlock>`	`0x1e824640`	prev/next coalesce lookup（同一比较器）
`GetBlockIf`	`0x1e818f20`	邻居物化；不存在/状态不符时为 null-slot Block

对齐量子 {#the-alignment-quantum}

对齐只在请求边界强制一次，之后仅由 block 边界追踪。

Allocate 时 round-up。 aligned = (size + align − (size != 0)) & −align，其中 align = alignment_in_bytes_（+0x60），是 2 的幂。(size != 0) 项让零字节请求保持为零，而不是向上取整到一个量子。
Region end 在构造时向下取整。 ctor 设置 allocatable_range_end_（+0x70）= end − (end % align)。结合 base_offset 是量子的倍数，这保证每个 offset、size 和 remainder 都是 alignment 的倍数，free list 永远不会产生亚量子 fragment。
上界。 max_aligned_size_（+0x68）= INT64_MAX − (INT64_MAX % align)；任何超过它的请求都会被 BadSizeError 拒绝，防止 round-up 算术溢出。

构造函数用 fatal LogMessageFatal 检查验证 Config，全部在反编译中确认：

text

base_offset_in_bytes_ >= 0
alignment_in_bytes_ > 0
allocatable_range_end_ > 0
alignment_in_bytes_ % granule_in_bytes_ == 0
(alignment_in_bytes_ & (alignment_in_bytes_ - 1)) == 0   "alignment_in_bytes_ must be a power of two"
```text

### 每层级的量子是什么

量子是该层级的 `Config.alignment_in_bytes_`，并且必须是 2 的幂、也是硬件 granule 的倍数。对 **HBM** 而言，它是 **1024 B DMA 最小值**（`jf_driver::kHbmMinimumDmaAlignment`）。DMA-issue site 会独立强制同一常量（`(address & (kHbmMinimumDmaAlignment − 1)) == 0`），因此使用更小 HBM 量子的分配器会发出被 DMA engine 拒绝的地址。完整 DMA alignment contract、16 KiB *编译期* program alignment（`xla_jf_program_hbm_alignment_in_kib`，在 MSA 前应用，因此静态 offset 到达时已预对齐）以及按层级表位于 [hbm-dma-alignment.md](hbm-dma-alignment.md)。本页只需要这一事实：分配器的量子*就是*该层级的 `Config.alignment_in_bytes_`，其下游所有内容都是 alignment 的倍数。

> **注意 —** 因为 round-up padding 归 allocation 所有（包含在 `HashTableEntry.size` 中），所以没有单独的 padding 记账。`size` 和 `aligned` 之间的字节只是 allocated block 的一部分，并会在 `Deallocate` 时归还给 free list。Compaction 的 `Memmove` size 以 *granule* 为单位表达（`size / granule_in_bytes_`），不是字节单位；见 [on-device-compaction.md](on-device-compaction.md)。

---

## 单池、固定 Region

每个 `(core, memory-tier)` 恰好有**一个** `BestFitAllocator`。设备侧没有 arena、没有 slab、没有 doubling 增长；region 是一个在启动时由 `Config` 建立的固定 `[base_offset_in_bytes_, allocatable_range_end_)` span。容量固定；`Expand`（`0x1e81a7a0`）和 `Shrink`（`0x1e81ada0`）会调整 region bounds，但它们由 program-stack manager 使用，而不是由通用分配使用。

整个内存栈中唯一类似 arena 的层是**主机** `PremappedMemoryManager`，它 round-robin N 个 2 的幂分区，每个分区在 `absl::Mutex` 下包装自己的 `BestFitAllocator`，这是主机 DMA staging 池，不是设备 HBM 分配器。它以及主机侧 `tsl::BFCAllocator`（只用于 XLA spill 到 host RAM 的 HBM buffer）在 [overview.md](overview.md) 和 [embedded-tcmalloc.md](embedded-tcmalloc.md) 中介绍。

> **警告 —** 单个 `BestFitAllocator` 实例**内部没有同步**。HAL 在其上层持有相关锁（主机 `PremappedMemoryManager` 用每分区 `absl::Mutex` 包装每个分区的分配器；设备侧 per-core HAL 序列化访问）。重新实现者不得假设 `Allocate`/`Deallocate` 自身是线程安全的。

---

## OOM 和碎片整理（指针）

当 `FindAllocatableBlock` 返回 `end()` 时，`Allocate` 返回 `ResourceExhausted` 状态：“…That was not possible. There are &lt;N&gt; free. The largest contiguous region of free memory is &lt;M&gt; due to fragmentation.”。（表面相似的 “…at the bottom of memory…” diagnostic 属于 `ReserveBottomOfMemory`，不是这条路径。）运行时的响应，即有界（≤ 2 次）重试、program eviction，以及发出 `TpuMemmoves` relocation plan 并把 block map + free tree 改写到 compaction 后布局的 `Compact` defrag，**归 [on-device-compaction.md](on-device-compaction.md) 所有**。属于本页的一个事实是：`Compact`（`0x1e81c360`）接受一个 **pinned** 地址的 `flat_hash_set<int64_t>`（MSA-static 和 DMA/aliasing-held buffer），这些地址绝不能被重定位，并会改写*此分配器的* `blocks_by_offset_` 和 `free_tree_` 到新布局，使 live `tpu::TpuBuffer` handle 自动解析到新 offset。pinned set 是强制“MSA 拥有静态布局；运行时拥有剩余空间中的动态布局”的精确机制。

---

## 与 MSA 的关系（指针）

[MSA](../compiler/msa-overview.md) 是*编译期* planner：它为每个 `HloValue` 选择 memory space，并（通过 `xla::jellyfish::ProgramMemoryAllocator`）选择一个冻结 offset，序列化进编译程序中的 `ProgramMemoryMetadata` proto。加载时，`ProgramMemoryAllocator::CreateFromProto` 重放这些 offset，而这个 `BestFitAllocator` 被*告知*每个静态 buffer 放在哪里，它不会为它们运行 best-fit 搜索。本页记录的 free-list 数学只运行于**动态**表面。这就是为什么运行时分配器和编译期 placement 是两个不同页面：概念问题相同（每个 buffer 在 HBM 中放在哪里），但代码路径完全不同，运行时间也完全不同。

---

## 交叉引用

- [overview.md](overview.md) — 五个片上层级 + 主机内存；本分配器在图中的位置
- [hbm-dma-alignment.md](hbm-dma-alignment.md) — 1024 B DMA-issue alignment contract，以及馈入本分配器量子的 16 KiB 编译期 program alignment
- [on-device-compaction.md](on-device-compaction.md) — `Compact`、`TpuMemmoves` relocation plan、OOM retry chain，以及通向 MSA 的 pinned-set 桥接
- [vmem-allocator.md](vmem-allocator.md) — 兄弟 VMEM 层级，同一个 `BestFitAllocator` 类，不同 `Config`
- [cmem-pool.md](cmem-pool.md) — CMEM 层级（`CmemSize == 0` 时休眠）
- [tpu-buffer-layout.md](tpu-buffer-layout.md) — 设备 buffer 的 offset 如何映射到设备端 tile layout
- [buffer-donation-aliasing.md](buffer-donation-aliasing.md) — compaction 期间变成 pinned 的 aliased buffer
- [embedded-tcmalloc.md](embedded-tcmalloc.md) — 主机侧 `tsl::BFCAllocator` / tcmalloc 路径（不同于此设备分配器）
- [../compiler/msa-overview.md](../compiler/msa-overview.md) — 编译期 HBM/VMEM placement pass，冻结本分配器重放的静态 offset
- [../compiler/msa-reservation-hbm-policy.md](../compiler/msa-reservation-hbm-policy.md) — 运行时分配器上游的按空间 capacity / reservation policy
- [back to index](../index.md)

HBM BestFit 分配器 ​

摘要 ​

范围与边界 ​

两个耦合的数据结构 ​

reserved-bottom watermark ​

Allocation 与 Split / Fragmentation 策略 ​

内部与外部碎片 ​

Free 时合并 ​

对齐量子 {#the-alignment-quantum} ​