PJRT 客户端分配器集成

本页中的所有地址、vtable 槽位和结构体大小都适用于 libtpu-0.0.40-cp314 wheel 中的 libtpu.so（build-id 89edbbe81c5b328a958fe628a9f2207d，build libtpu_lts_20260413_b_RC00，clang 9999.0.0）。其他版本会有所不同。

摘要

本页记录 TPU PJRT 插件内部的 stream-executor 分配器桥接层：它位于 PJRT 的设备内存 API 与核上 best-fit 引擎之间。在 StreamExecutor 后端中，se::DeviceMemoryAllocator::Allocate(ordinal, size, …) 调用会返回一个 se::OwningDeviceMemory（一个 ScopedDeviceMemory），其中包装了 DeviceMemoryBase。libtpu 在热路径上没有 se::StreamExecutorMemoryAllocator，而是提供了一条承担相同角色的定制链路。xla::TpuClient::AllocateRawBuffer 是 Allocate 入口；TpuSharedMemoryLocation 是 DeviceMemoryBase 身份（ordinal + memory-space 层级）；tpu::TpuBuffer（由 xla::TpuRawBuffer 包装）是拥有所有权的句柄，其析构函数负责释放，扮演 ScopedDeviceMemory 的角色；tpu::System::Allocate 是按 ordinal 路由的分派器，会分派到正确核心的分配器。

这个桥接层是 StreamExecutor 读者必须重新映射其心智模型的部分。AllocateRawBuffer 按 PjRtMemorySpace::kKindId（HBM 设备、pinned host、unpinned-host/CPU staging）路由到三个后端之一，这相当于一个多 memory-space 的 Allocate。设备分支会穿过四层间接调用：TpuClient::AllocateBuffer → tpu::AllocateBuffer → tpu::System::Allocate（按核心 map 查找，即 ordinal → allocator 步骤）→ 虚函数 tpu::TpuAllocator::Allocate → tpu::TpuSharedMemory::AllocateLocked → HBM best-fit Allocate。Deallocate 侧是对称的，从 TpuBuffer 析构函数出发，通过同一个 TpuAllocator vtable 返回。本页负责说明这条链路、Allocate/Deallocate ABI 以及 memory-space 路由；它有意停在 best-fit Allocate 调用边界。

注意 — best-fit-with-coalescing 算法（free-list 结构、RB-tree 搜索、积极合并、拆分策略、1024 B HBM 量子）位于 ../memory/hbm-allocator.md。本页调用 BestFitAllocator::Allocate（vt+0x30，0x1e817820）和 Deallocate（vt+0x38，0x1e819dc0）后即停止，不会重新推导该引擎的内部行为。

对重新实现而言，桥接契约是：

Allocate ABI。 TpuClient::AllocateRawBuffer(PjRtMemorySpace*, size, bool, AsyncValueRef<bool>) 按 memory-space kind 路由。设备分支恰好经过四层间接调用到达引擎，最后在槽位 +0x10 处进行 vtable tail-call。
按 ordinal 路由器。 tpu::System::Allocate 查找 flat_hash_map<TpuSharedMemoryLocation, unique_ptr<TpuAllocator>>，并 tail-call 匹配分配器的虚拟 Allocate。TpuSharedMemoryLocation 是等价于 DeviceMemoryBase 的键。
TpuAllocator 6 槽虚接口。 两种具体策略（ReusingTpuAllocator、DeferredTpuAllocator）共享同一种 vtable 形状；两者都把物理放置委托给按核心的 TpuSharedMemory。
ScopedDeviceMemory 所有权。 tpu::TpuBuffer（变体 {Owned, Sliced, Unsafe}）是拥有所有权的句柄；xla::TpuRawBuffer 为 PJRT 包装它；析构通过 TpuAllocator::Deallocate（vt+0x30）回路由。
调用落点以及停止位置。 TpuSharedMemory::AllocateLocked 分派到 BestFitAllocator::Allocate（vt+0x30）；引擎内部见 ../memory/hbm-allocator.md。


`Allocate` 入口（PJRT）	`xla::TpuClient::AllocateRawBuffer` `0xf7fb1e0` — 按 `PjRtMemorySpace` kind 路由
设备子入口	`xla::TpuClient::AllocateBuffer` `0xf7fc5a0` → `tpu::AllocateBuffer` `0xf8d51c0`
按 ordinal 路由器	`tpu::System::Allocate(loc, size)` `0x1d0aeea0` — `flat_hash_map` 查找 → `jmp *vt[+0x10]`
分配器键（`DeviceMemoryBase`）	`tpu::TpuSharedMemoryLocation` — ctor `0x20ad6ae0`，`operator==` `0x20ad6be0`，hash `0x1d0ba2a0`
`TpuAllocator` 接口	抽象基类，typeinfo `0x21ca8508`；6 槽 vtable；`Allocate` = vt+0x10
具体策略	`ReusingTpuAllocator`（56 B，vt `0x21ca85d8`）· `DeferredTpuAllocator`（176 B，vt `0x21ca84a8`）
拥有所有权的句柄（`ScopedDeviceMemory`）	`tpu::TpuBuffer`（由 `xla::TpuRawBuffer` 包装）；设备端大小在 `TpuBuffer+0x50`
引擎边界（不属于本页）	`tpu::BestFitAllocator::Allocate` vt+0x30 `0x1e817820` → `../memory/hbm-allocator.md`
置信度	高（字节锚定），除非某行或标注另有说明

范围和边界

本页是 stream-executor 分配器桥接层，负责 Allocate/Deallocate ABI 以及进入按核心分配器的 memory-space 路由。相邻主题位于各自页面；不要在这里重复。

关注点	所属页面
best-fit free-list 结构、RB-tree 搜索、合并、拆分策略、对齐量子	`../memory/hbm-allocator.md`
`tpu::TpuBuffer` 字段布局、`TpuSharedMemoryLocation` 编码、slice/view 语义	`../memory/tpu-buffer-layout.md`
buffer donation、input/output aliasing、`ScopedHold`、`HloInputOutputAliasConfig`	`../memory/buffer-donation-aliasing.md`
`ExecuteAsyncOnStream` 调用输出分配的位置、dispatch/launch 流程	`execute-async-on-stream.md`
程序加载、scratch 绑定、in-flight-execution semaphore	`load-program-enqueue.md`

本页负责：AllocateRawBuffer → … → AllocateLocked 调用链、Allocate/Deallocate vtable ABI、TpuSharedMemoryLocation ordinal/tier 路由，以及把 TpuBuffer 作为 ScopedDeviceMemory 的所有权接线。ExecuteAsyncOnStream 执行的输出缓冲区分配只在它进入该桥接层的位置展示。

StreamExecutor 心智映射

来自 XLA StreamExecutor 后端的读者应始终记住以下对应关系。libtpu 不会为 runtime buffer 实例化 se::StreamExecutorMemoryAllocator；它用 TPU 专用类型重新实现相同职责。

StreamExecutor 概念	libtpu 桥接对象
`se::DeviceMemoryAllocator::Allocate(ordinal, size, …)`	`xla::TpuClient::AllocateRawBuffer`（按 memory-space 路由）
`Allocate` 设备 ordinal 选择	`tpu::System::Allocate` map 查找 `loc → allocator`
`se::DeviceMemoryBase`（不透明 base + size）	`tpu::TpuSharedMemoryLocation`（chip、tier、host index）
`se::OwningDeviceMemory` / `ScopedDeviceMemory`（RAII）	`tpu::TpuBuffer`（`xla::TpuRawBuffer` PJRT 包装器）
`se::DeviceMemoryAllocator::Deallocate`	`tpu::TpuAllocator::Deallocate`（vt+0x30）
具体后备分配器	按核心的 `tpu::TpuAllocator` → `tpu::BestFitAllocator`

特性 — DeviceMemoryBase 类比物（TpuSharedMemoryLocation）不是一个 raw pointer 加 size。它是结构化的身份：(chip, on-chip-memory-segment, host index)，同时作为选择哪个核心分配器服务请求的 hash-map key（operator== 0x20ad6be0，AbslHashValue 0x1d0ba2a0）。“ordinal”和“memory space”融合在同一个键中。把 ordinal 建模为裸整数的重新实现会错误路由 VMEM/SMEM/CMEM 层级，因为 tier 是同一个 key 的一部分，而不是单独参数。

`Allocate` 路径（PJRT_Buffer → Best-Fit 引擎）

目的

这是桥接层的正向路径：PJRT buffer-create、ExecuteAsyncOnStream 输出缓冲区，或 transfer-staging 请求从 AllocateRawBuffer 进入，按 memory space 路由，并且（对设备层级而言）下行到按核心的 best-fit Allocate。每一跳都是薄适配器；之所以存在间接层，是为了让一个进程级 System 能为每个核心池化一个分配器，并在所有 PJRT client 之间共享。

入口点

text

xla::TpuClient::AllocateRawBuffer (0xf7fb1e0)        ── Allocate ABI; routes by PjRtMemorySpace kind
  ├─ [HBM]  TpuClient::AllocateBuffer (0xf7fc5a0)    ── builds TpuSharedMemoryLocation (the ordinal/tier key)
  │    └─ tpu::AllocateBuffer (0xf8d51c0)            ── sync vs AllocateAfter; TraceMe; latency log
  │         └─ tpu::System::Allocate (0x1d0aeea0)    ── per-core map lookup → jmp *vt[+0x10]
  │              └─ TpuAllocator::Allocate (virtual) ── Reusing (0x1d0d2480) | Deferred (0x1d0ce900)
  │                   └─ TpuSharedMemory::AllocateLocked (0x1d4be920 / inner 0x1d4c0f40)
  │                        └─ BestFitAllocator::Allocate vt+0x30 (0x1e817820)  ── ENGINE (other page)
  ├─ [PINNED HOST]   tpu::System::AllocateHostBuffer (0x1d0af180)  ── pinned host (separate engine)
  └─ [UNPINNED HOST] xla::CpuRawBuffer::Allocate (0xf911680)       ── CPU-resident staging
```text

### 算法

路由器和按 ordinal 查找是重新实现者必须精确复现的两个决策。链路其余部分都是转发。

```c
function TpuClient_AllocateRawBuffer(memspace, size, is_tuple, dep):  // 0xf7fb1e0
    // The Allocate-ABI entry. Routing is by PjRtMemorySpace::kKindId, the
    // multi-memory-space equivalent of se::DeviceMemoryAllocator::Allocate.
    // `dep` is the AsyncValueRef<bool> allocate-after gate.
    // Tested in this order (no default device branch — HBM is the fallthrough):
    if memspace.kind() == UnpinnedHostMemorySpace::kKindId:   // CPU-resident staging
        CHECK(!dep)                               // LogFatal: allocate_after unsupported (unpinned host)
        return CpuRawBuffer_Allocate(memspace, size,
                                     CpuDeviceMemory_DefaultAllocator())  // 0xf911680
    if memspace.kind() != TpuHbmMemorySpace::kKindId:
        if memspace.kind() == PinnedHostMemorySpace::kKindId: // pinned host staging
            CHECK(!dep)                           // LogFatal: allocate_after unsupported (pinned host)
            return System_AllocateHostBuffer(loc, size, deleter)         // 0x1d0af180
        return error("Unsupported memory space: %s.")
    // else: TpuHbmMemorySpace  -> DEVICE
    return TpuClient_AllocateBuffer(memspace.device(), size, is_tuple, dep) // 0xf7fc5a0

function TpuClient_AllocateBuffer(dev, size, is_tuple, dep):           // 0xf7fc5a0
    // Build the DeviceMemoryBase-equivalent key: (chip, tier, host index).
    loc = dev.LocalSharedMemory(kHbm).index_on_host()  // 0x20ad6840 / 0x20ad6e00
    pel = PendingEventLoggers_get(loc.index())         // 0x1d0a89a0
    jitter = SharedBitGen()()                          // 0xf822540  (retry/backoff seed)
    return tpu_AllocateBuffer(loc, size, /*sync*/, /*async*/,
                              dep, System*, work_queue, pel)            // 0xf8d51c0
    // on ResourceExhausted: LogOnResourceExhausted (0xf7fc4c0) then
    //   error::RuntimeBufferAllocationFailure (0xf7fd3a0)

function System_Allocate(this, loc, size):             // 0x1d0aeea0  (byte-confirmed)
    // The per-ordinal router == se::DeviceMemoryAllocator's ordinal step.
    // map type: flat_hash_map<TpuSharedMemoryLocation, unique_ptr<TpuAllocator>>
    alloc = map.at(this->allocators_, loc)             // raw_hash_map::at(*loc+48, ...)
    // tail-call the matched allocator's virtual Allocate (vtable slot +0x10):
    return alloc->vtable[+0x10](this, alloc, size)     // jmp *0x10(%rax)

注意 — tpu::AllocateBuffer（0xf8d51c0）会根据是否存在 AsyncValueRef<bool> 依赖来选择 System::Allocate（同步，0x1d0aeea0）或 System::AllocateAfter（0x1d0af060）以门控分配，也就是 tpu_allow_async_allocations 配置路径。AllocateAfter 将分配排在进行中的 free 之后；两个 bool 参数选择 sync-vs-after 和 async 开关。vtable 还在槽位 +0x18 携带 AllocateAfter，因此 async 变体也通过同一个按核心分配器路由。

反编译交叉检查 — 路由器 tail-call

核心结论是 System::Allocate 是一次 map 查找，随后对槽位 +0x10 进行 vtable tail-call；这一点由字节确认：

// _ZN3tpu6System8AllocateERKNS_23TpuSharedMemoryLocationEl @ 0x1d0aeea0
__int64 tpu::System::Allocate(System *this, const TpuSharedMemoryLocation *a2,
                              __int64 size, __int64 a4) {
  _QWORD *v5 = raw_hash_map<FlatHashMapPolicy<TpuSharedMemoryLocation,
                 unique_ptr<TpuAllocator>>>::at(*(_QWORD*)a2 + 48LL, size);
  //          map handle lives at *loc+48; .at() returns &unique_ptr<TpuAllocator>
  return (*(__int64(**)(System*, _QWORD, __int64))
            (*(_QWORD*)*v5 + 16LL))(this, *v5, a4);   // jmp *vt[+0x10] = TpuAllocator::Allocate
}
```text

`TpuAllocator` vtable 的槽位 `+0x10` 是 `Allocate`；该函数以解包后的 `unique_ptr` payload（`*v5`）作为 `this` 来 tail-call 它。这是路由器与策略之间的 `Allocate` ABI 边界。

### 函数映射

| 函数 | 地址 | 角色 |
|---|---|---|
| `xla::TpuClient::AllocateRawBuffer` | `0xf7fb1e0` | `Allocate` ABI 入口；memory-space 路由 |
| `xla::TpuClient::AllocateBuffer` | `0xf7fc5a0` | 设备分支；构造 `TpuSharedMemoryLocation` |
| `tpu::AllocateBuffer` | `0xf8d51c0` | sync/`AllocateAfter` 选择；trace + latency |
| `tpu::System::Allocate` | `0x1d0aeea0` | 按 ordinal map 查找 → vt+0x10 tail-call |
| `tpu::System::AllocateAfter` | `0x1d0af060` | Async-gated 变体（vt+0x18） |
| `tpu::TpuSharedMemory::AllocateLocked` | `0x1d4be920` / `0x1d4c0f40` | 分派 `BestFitAllocator::Allocate`（vt+0x30） |
| `xla::CpuRawBuffer::Allocate` | `0xf911680` | CPU-resident memory-space 分支 |

---

## `TpuAllocator` 虚接口（`Allocate`/`Deallocate` ABI）

### 目的

`tpu::TpuAllocator` 是一个抽象基类，其 6 槽 vtable *就是*按 ordinal 路由器与引擎之间的桥接 ABI。它呈现“FFI 形状”：一个稳定虚表，路由器在不知道给定核心背后是哪种具体策略的情况下进行 tail-call。两种策略实现它；两者最终都把物理放置交给按核心的 `TpuSharedMemory`（进而交给 `BestFitAllocator`）。

### vtable

从 `0x21ca84a8` 处的 `DeferredTpuAllocator` vtable 恢复得到（`0x21ca85d8` 处的 `ReusingTpuAllocator` 镜像具有相同形状）：

| 槽位 | 方法 | Deferred 实现 | Reusing 实现 |
|---|---|---|---|
| `vt[+0x00]` | `~dtor`（deleting） | `0x1d0ce4c0` | — |
| `vt[+0x08]` | `~dtor`（complete） | `0x1d0ce8c0` | — |
| `vt[+0x10]` | `Allocate(long)` | `0x1d0ce900` | `0x1d0d2480` |
| `vt[+0x18]` | `AllocateAfter(long, RCRef<AsyncValue>)` | `0x1d0ce920` | `0x213ba520` |
| `vt[+0x20]` | `RegisterSequencedOutOfMemoryHold(RCRef<AsyncValue>)` | `0x1d0cec80` | `0x1d0d28c0` |
| `vt[+0x28]` | `Shutdown(bool)` | `0x1d0cf360` | `0x1d0d28e0` |
| `vt[+0x30]` | `Deallocate(TpuBuffer*, InlinedVector<RCRef<AsyncValue>,2>)` | `0x1d0cf460` | `0x1d0d2c00` |

> **特性 —** `Allocate` 位于 vt+0x10，但 `Deallocate` 位于 vt+**0x30**，并不相邻。中间两个槽位是 `AllocateAfter` 以及 OOM-hold/`Shutdown` 对。假设 `Allocate`/`Deallocate` 在 vtable 中相邻的重新实现（许多 `se::DeviceMemoryAllocator` 布局确实如此）会错误分派 free。free 侧携带一个包含 completion event 的 `InlinedVector<RCReference<AsyncValue>,2>`，因此 deallocation 可以排在进行中的使用之后，这是 deferred-deallocation 钩子。

### 两种策略

```c
function TpuAllocator_Create(shared_mem, strategy, tracker):  // 0x1d0ce420 (byte-confirmed)
    if strategy == 2:                          // kReusing
        a = operator new(0x38)                 //  56-byte instance
        a.vptr      = ReusingTpuAllocator_vtable + 0x10   // off_21CA85E8
        a[+0x08]    = shared_mem                // the per-core TpuSharedMemory
        a[+0x10..]  = { 0, 1, 0 }              // reuse-cache bookkeeping
    else:                                      // kDeferred (default)
        a = operator new(0xB0)                 // 176-byte instance
        a.vptr      = DeferredTpuAllocator_vtable + 0x10   // off_21CA84B8
        a[+0x08]    = shared_mem
        zero a[+0x10 .. +0xA0]                 // tracker/bitset state
        a[+0xA8]    = tracker                  // SystemEventTracker* (result[21])
    return a

ReusingTpuAllocator::Allocate（0x1d0d2480）→ AllocatorFor(size)（0x1d0d2e80）：在共享 mutex 下选择按 size-class 的子分配器 / reuse cache；如果有合适的已释放回收 buffer，则返回它，否则新分配。这是 size-class reuse 策略，最接近 caching se allocator 的类比物。
DeferredTpuAllocator::Allocate（0x1d0ce900）→ AllocateImpl(size)（0x1d0cf620）：已验证函数体是一行转发到 AllocateImpl，随后 return this。AllocateImpl 获取 mutex（对 this+0x10 调用 absl::Mutex::lock）、跟踪分配，并用 MakeErrorImpl<13> 防护缺失核心（absl::StatusCode::kInternal，消息 "No attached TPU to allocate with."，字符串锚定到 tpu_allocator.cc）；物理放置通过间接 vtable 调用委托到所属 TpuSharedMemory/driver 对象，实际 no-room ResourceExhausted/fragmentation 诊断在下面的 AllocateLocked 中产生。Deferred free 会批处理并回收，这就是“deferred deallocation”名称的来源。

// _ZN3tpu12_GLOBAL__N_120DeferredTpuAllocator8AllocateEl @ 0x1d0ce900 (byte-confirmed)
DeferredTpuAllocator *Allocate(DeferredTpuAllocator *this, __int64 size, __int64 a3) {
    DeferredTpuAllocator::AllocateImpl(this, size, a3);   // 0x1d0cf620
    return this;
}
```text

> **注意 —** `TpuAllocator::Create`（`0x1d0ce420`）由字节确认：`strategy == 2` 选择 56 字节的 `ReusingTpuAllocator`，其他值选择 176 字节的 `DeferredTpuAllocator`，`TpuSharedMemory*` 存储在 `+0x08`，且（仅 deferred）`SystemEventTracker*` 存储在 `+0xA8`。除此三个槽位外，具体字段布局（reuse-cache map、per-allocation tracker bitset）未逐字段解码（`+0x08`/`+0xA8` 之后字段含义置信度低）。

### 调用落点（以及停止位置）

`TpuSharedMemory::AllocateLocked`（`0x1d4be920`，inner `0x1d4c0f40`）是本页负责的最后一跳。它受 mutex 保护，为 OOM 诊断获取引擎统计信息（`BytesAllocated` vt+0x18，`BytesReserved` vt+0x20，`BytesAvailable` vt+0x68，`BytesAllocatable` vt+0x78，`GetFragmentation` vt+0x80），然后分派：

```c
function TpuSharedMemory_AllocateLocked(this, size):   // 0x1d4be920 / inner 0x1d4c0f40
    engine = this->best_fit_                            // per-tier BestFitAllocator
    offset = engine->vtable[+0x30](engine, aligned_size) // call *0x30(%rax) -> 0x1e817820
    //  ^^^ BEST-FIT ALLOCATE — algorithm owned by ../memory/hbm-allocator.md
    return TpuBuffer(loc, offset, ...)                  // wrap offset into the owning handle

反编译确认 AllocateLocked 函数体中存在 + 48（0x30）分派，即 (**(this+192) + 48LL)(…)，其中 this+192 是按 tier 的 best_fit_ 引擎指针（0x1d4c0f40 处的 inner $_0 是 OOM-stats 诊断构造器，会读取 BytesReserved vt+0x20、BytesAllocated vt+0x18、BytesAvailable vt+0x68、BytesAllocatable vt+0x78、GetFragmentation vt+0x80）。+0x30 调用之后的一切，包括 RB-tree best-fit 搜索、SplitBlock、合并、对齐向上取整，都在 ../memory/hbm-allocator.md。

`Deallocate` 路径（`ScopedDeviceMemory` 释放）

目的

tpu::TpuBuffer 是拥有所有权的句柄，即 se::OwningDeviceMemory / ScopedDeviceMemory 的类比物。当最后一个引用释放时，free 会通过与分配进入时相同的 TpuAllocator vtable 返回，位于槽位 +0x30。这里没有单独的 se::DeviceMemoryAllocator::Deallocate(ordinal, mem)；buffer 携带自己的 allocator() 和 location()，因此 free 可以自路由。

算法

function ~TpuBuffer():                                  // owning-handle destructor (Owned variant)
    if variant == Sliced or variant == Unsafe: return   // views/unsafe do not own
    alloc = this->allocator()                           // 0x1d102680  -> TpuAllocator*
    loc   = this->location()                            // 0x1d102740  -> TpuSharedMemoryLocation
    alloc->vtable[+0x30](alloc, this, pending_events)   // TpuAllocator::Deallocate (vt+0x30)
        // Deferred: 0x1d0cf460 — may batch the free, gated by the deferred-dealloc flags
        // Reusing:  0x1d0d2c00 — returns the block to its size-class reuse cache
            // -> TpuSharedMemory::DeallocateLocked -> BestFitAllocator::Deallocate (vt+0x38)
            //    coalesce-on-free is the engine's job (other page)
```text

> **易错点 —** `Owned`/`Sliced`/`Unsafe` 变体标签（`TpuBufferTraits`）决定析构函数是否释放。`Sliced` buffer 是共享父分配的子 buffer *view*（`Slice` `0x1d1017c0` / `0x1d101f60` / `0x1d102f60`）；`Unsafe` buffer（`IsUnsafe` `0x1d102800`，`UnsafeReleaseBuffer` `0x1d102660`）已经释放所有权。只有 `Owned` 会 free。对每个析构函数都 free 的重新实现会对每个 sliced view 造成 double-free。变体模型见 [`../memory/tpu-buffer-layout.md`](../memory/tpu-buffer-layout.md)。

Deferred deallocation 是唯一偏离普通 `ScopedDeviceMemory` 的行为：它不是立即释放，`DeferredTpuAllocator::Deallocate`（`0x1d0cf460`）可以将 free 入队，并受 `FLAGS_tpu_deferred_deallocation`（`0x22396ee0`）/ `FLAGS_tpu_reap_deferred_deallocations`（`0x22396f40`）/ `…_try_wait`（`0x22397008`）门控。随后 pending free 可通过 `RegisterSequencedOutOfMemoryHold`（vt+0x20）满足受阻分配，也就是链接在 `AsyncValue` 上的“sequenced” hold。

---

## Memory-Space 和 Ordinal 路由

### 目的

该桥接层的路由有两个轴，而 StreamExecutor 读者常会把它们合并为一个：**memory-space kind**（选择哪个后端）和 **ordinal/tier**（哪个核心、哪个 on-chip tier）。`AllocateRawBuffer` 决定前者；`System::Allocate` 的 map key 决定后者。

### 轴 1 — memory-space kind（后端 switch）

| `PjRtMemorySpace::kKindId` | 后端 | 引擎 | 本页？ |
|---|---|---|---|
| `TpuHbmMemorySpace`（fallthrough） | `TpuClient::AllocateBuffer` → `System::Allocate` | `tpu::BestFitAllocator` | 负责桥接层 |
| `PinnedHostMemorySpace` | `tpu::System::AllocateHostBuffer`（`0x1d0af180`） | premapped pool / `HostBufferPool` | 相邻 |
| `UnpinnedHostMemorySpace` | `xla::CpuRawBuffer::Allocate`（`0xf911680`） | `CpuDeviceMemory::DefaultAllocator` | 相邻 |
| 其他任何值 | `MakeRep` + `"Unsupported memory space: %s."` | — | 错误 |

设备分支是本页端到端追踪的分支。这里列出 host 和 CPU-staging 分支，是因为 `AllocateRawBuffer` 是共享的 `Allocate` 入口，会按 `kKindId` 在它们之间选择；它们的引擎属于独立子系统。`allocate_after` async 依赖会对两类 host 都被拒绝（`LogFatal`）；只有 HBM 设备分支携带它。

### 轴 2 — `TpuSharedMemoryLocation` 键（ordinal/tier）

路由器的 map 是 `flat_hash_map<TpuSharedMemoryLocation, unique_ptr<TpuAllocator>>`，也就是**每个核心一个 `TpuAllocator`**。键编码 `(chip, on-chip-memory-segment, host index)`：

```c
struct tpu::TpuSharedMemoryLocation {   // the DeviceMemoryBase identity
    // ctor 0x20ad6ae0: (TpuTopology*, TpuDimensions, TpuDimensions, TpuSharedMemoryOnChip)
    // accessors: Chip(), index_on_host() (0x20ad6e00), ToString() (0x20ad6f00)
    // operator== 0x20ad6be0 ; AbslHashValue 0x1d0ba2a0  -> these make it a hash key
};

由于 tier（TpuSharedMemoryOnChip）在 key 内部，HBM、VMEM、SMEM、CMEM 和 SFLAG 各自映射到不同的 TpuAllocator/BestFitAllocator 实例。同一条 System::Allocate 代码路径服务所有 tier，区别只在 key。因此 ../memory/hbm-allocator.md 可以说同一个 BestFitAllocator 类支撑每个 tier：桥接层在引擎运行前就已经按 key 选择了按 tier 的实例。

`ExecuteAsyncOnStream` 如何进入该桥接层

ExecuteAsyncOnStream 的输出缓冲区分配是该桥接层的主要 runtime 调用方。完整 dispatch/launch 流程见 execute-async-on-stream.md；这里仅展示入口位置。

对于每个输出 ShapeIndex，TPU 输出分配器会决定 reuse-vs-fresh，而 fresh allocation 会直接调用该桥接层：

text

ExecuteAsyncOnStream  (execute-async-on-stream.md)
  └─ tfrt::tpu::AllocateOutputBuffersWithInputReuse (0xf7ba9a0)
       ├─ aliased output  → REUSE the donated input's TpuBuffer in place   (no allocation)
       │                    (HloInputOutputAliasConfig::GetAliasedParameter 0x1e580200)
       │                    → buffer-donation-aliasing.md
       └─ non-aliased     → tfrt::tpu::AllocateTpuBufferWithRetry (0xf7ec6a0)
                              └─ tpu::System::Allocate (0x1d0aeea0)  ── THIS BRIDGE
```text

> **注意 —** donation/aliasing 决策（哪些输出复用哪些输入、`ScopedHold` pin、`HloInputOutputAliasConfig`）属于 [`../memory/buffer-donation-aliasing.md`](../memory/buffer-donation-aliasing.md)。本页只关心*非 alias* 分支，即通过 `System::Allocate` 执行真实分配的分支。alias 分支**不会**执行分配；它把 donated input 的 `TpuBuffer` 重新绑定为输出句柄。

`AllocateTpuBufferWithRetry`（`0xf7ec6a0` → `0xf7ed620` → `0xf7ed980` → `0xf7edd80`，一条递归链）用 defragment-and-retry 包装 `System::Allocate`：在 `ResourceExhausted` 时调用 `tpu::System::CompactMemory`（`0x1d0b6000`），然后重试。确切重试次数和 `SharedBitGen` backoff jitter 未经字节追踪（低）。最终 PJRT 可见错误是 `xla::error::RuntimeBufferAllocationFailure`（`0xf7fd3a0`）；叶子消息是 [`../memory/hbm-allocator.md`](../memory/hbm-allocator.md) 上的 `BestFitAllocator` fragmentation 诊断。

---

## 分配器生命周期（一个引擎，所有 client）

该桥接层的分配器**不是按 PJRT client 创建**的。它们在进程级 singleton 中按核心池化并共享：

```text
xla::GetSingletonTpuStatesManager (0xf958360, mutex + __cxa_guard)
  └─ xla::TpuStatesManager (one per process)
       └─ GetOrCreateTpuSystemState (0xf956e40) → xla::TpuSystemState
            └─ xla::CreateTpuSystemState → the singleton tpu::System
                 └─ flat_hash_map<TpuSharedMemoryLocation, unique_ptr<TpuAllocator>>
                      (ONE TpuAllocator per core, shared across every TpuClient)

进程中的每个 xla::TpuClient 都共享同一个 TpuSystemState/System（由 use_global_tpu_system 配置门控，字符串 0xa2e9960）。因此设备分配器是按设备（核心）创建、在 System 层池化，并在所有 PJRT client 之间共享，这与按 se::StreamExecutor 创建分配器实例相反。

特性 — 这是与 StreamExecutor 模型最深的差异。se::StreamExecutorMemoryAllocator 通常由一个 client/executor 拥有；这里的 Allocate/Deallocate ABI 有意路由到共享的按核心分配器，因此指向同一核心的两个 PJRT client 会竞争同一个 free-list。每 client 实例化一个分配器的重新实现不会看到相同的 fragmentation 或 OOM 行为，因为真实实现会跨 client 池化。

重新实现注意事项

四跳间接调用是有意设计，不是偶然。 AllocateRawBuffer → AllocateBuffer → tpu::AllocateBuffer → System::Allocate 的存在，是为了让 PJRT 可见 ABI 保持稳定，同时分层加入按核心分配器、async-vs-sync 选择以及 trace/latency instrumentation。折叠它会丢失 AllocateAfter async 路径和按核心共享。
路由按 key 发生两次。 Memory-space kind 在 AllocateRawBuffer 选择后端；TpuSharedMemoryLocation 在 System::Allocate 选择按核心分配器。两者都是必需的；二者都不是普通 ordinal。
句柄拥有 free。 不同于 se::DeviceMemoryAllocator::Deallocate(ordinal, mem)，libtpu 的 TpuBuffer 携带 allocator() 和 location()，因此 free 通过 vt+0x30 自路由，但只对 Owned 变体如此。
不要重新推导引擎。 调用落在 BestFitAllocator::Allocate（vt+0x30）和 Deallocate（vt+0x38）；算法见 ../memory/hbm-allocator.md。本页停在分派处。

交叉引用

../memory/hbm-allocator.md — 该桥接层调用的 best-fit-with-coalescing 引擎（Allocate vt+0x30，Deallocate vt+0x38）；负责 free-list 算法
../memory/tpu-buffer-layout.md — tpu::TpuBuffer 字段布局、Owned/Sliced/Unsafe 变体、TpuSharedMemoryLocation 编码
../memory/buffer-donation-aliasing.md — ScopedHold、HloInputOutputAliasConfig、绕过分配的 aliased-output reuse 路径
execute-async-on-stream.md — runtime 调用方；输出缓冲区分配进入该桥接层的位置
load-program-enqueue.md — 程序加载、预规划的 static scratch 绑定，以及 in-flight-execution semaphore（不是设备内存 reservation）
overview.md — runtime/execution 子系统地图

PJRT 客户端分配器集成 ​

摘要 ​

范围和边界 ​

StreamExecutor 心智映射 ​

Allocate 路径（PJRT_Buffer → Best-Fit 引擎） ​

目的 ​

入口点 ​

反编译交叉检查 — 路由器 tail-call ​

Deallocate 路径（ScopedDeviceMemory 释放） ​