主机回调

本页所有地址均适用于来自 libtpu-0.0.40-cp314 wheel 的 libtpu.so（构建 libtpu_lts_20260413_b_RC00，build-id md5 89edbbe81c5b328a958fe628a9f2207d，781,691,048 字节）。该镜像未 stripped；反混淆后的 C++ 符号名按原样引用。VA == 文件分析地址。其他版本会不同。

摘要

主机回调是一个正在运行的 TPU 程序在执行中途回到主机代码的方式：HLO 编译器发出一对 Send/Recv（即 xla.host_compute / outside-compilation 的 lowering，以及跨主机 send/recv custom-call 的 lowering），运行时设备在 sync flag 上暂停，主机运行已注册的回调，主机发出完成信号并交回（或消费）一个数据块后，设备继续执行。XLA 参考框架是 SendCallback/RecvCallback 这个上游称为 PjRtClient::ExecuteOptions::send_callbacks / recv_callbacks 的表面：每个回调由 HLO Send/Recv op 中携带的整数 channel id 索引，而一个 rendezvous 会把某个 channel 的设备侧与为同一 channel id 注册的主机回调匹配起来。这本质上是一个按 channel 键控的 rendezvous，不是 FIFO —— 对比 Infeed / Outfeed 队列，其中一次传输命名一个 {TpuCoreLocation, int queue_index}，设备按程序顺序消费条目，没有逐次传输的身份标识。

和 infeed/outfeed 一样，libtpu 将这套机制发布了两份，重新实现者不能混淆二者。现代 PJRT 路径是 xla::TpuHostTransferManager（learning/45eac/research/pjrt/tpu_pjrt_client.cc）：每次执行一个 transfer manager，用 {TpuClient*, TpuCoreLocation} 对构造，持有两个按 channel 键控的回调 map —— send map（device→host，"CopyFromDevice"）和 recv map（host→device，"CopyToDevice"）—— 每个都是 absl::flat_hash_map<uint32_t channel, callback>。它由 TpuExecutableLoadState::ExecuteLaunchRaw 通过 TpuHostTransferManager::SetExecuteEvent(AsyncValueRef<tpu::TpuEvent>) 注册到执行路径上，并在两个专用的 detached 线程上排空 chunk（TpuHostTransferManagerSendThread / TpuHostTransferManagerRecvThread）。遗留路径是 tensorflow::TpuHostTransferManagerBase（learning/45eac/google/xla/tpu_host_transfer_manager.cc），由 xla::LocalClient / TF-TPU op kernel 使用：它把 rendezvous_key_base 字符串解析为 TF RendezvousInterface，保留一个 std::map<int channel, HostSendRecvInfo>，并由 TPU driver 触发 HandleHostCommand 驱动；该函数解码 32 位命令字 {high byte = Send|Recv, low 24 bits = channel}，并通过 rendezvous_->Send / rendezvous_->RecvAsync(transfer.parsed_key, …) 进行匹配。

本页负责主机回调的 rendezvous + transfer-manager 注册 + channel-id 键控。设备端 Send/Recv 的 sync-flag 机制最终落到 TPU driver（deepsea::executor::DeepseaStream::EnqueueOnDeviceSendRecvLocal / EnqueueOnDeviceSend / EnqueueOnDeviceRecv），这里只描述到契约深度。通用的 Stream::DoHostCallbackWithStatus 主机回调 shim（StreamExecutor 层级的 host_compute trampoline）是遗留路径也可抵达的叶子节点，本页也覆盖它；注册 transfer manager 的执行路径见 Execute Async on Stream；批量 infeed/outfeed FIFO 见 Infeed / Outfeed 队列。

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

两个 transfer manager —— PJRT xla::TpuHostTransferManager（每次 launch 一个，{TpuClient*, TpuCoreLocation}，两个 flat_hash_map<uint32_t,callback>）对比遗留 tensorflow::TpuHostTransferManagerBase（std::map<int,HostSendRecvInfo> + TF RendezvousInterface）—— 以及它们是相互独立的代码路径，而不是一个包装另一个。
channel-id 键 —— 回调通过 HLO Send/Recv 中的整数 channel id 查找。没有已注册回调的设备侧命令是致命的（LogMessageFatal），不是静默丢弃。PJRT 在 SwissTable 中以 uint32_t 为键；遗留路径在红黑树中以 int 为键。
注册 → 触发 → 继续的 rendezvous —— 注册发生在 launch 前（PJRT 上是 SetExecuteEvent，遗留路径上是 Initialize(rendezvous_key_base, RendezvousInterface*)）；设备遇到 Send/Recv 时触发主机侧；主机回调运行，通过 sync flag 发信号，设备继续执行。manager 持有的 execute event 完成整个 launch。
数据方向不对称 —— Send op 是 device→host（主机通过 HandleSendChunk 接收 PjRtChunk，即 CopyFromDeviceCallback）；Recv op 是 host→device（主机通过 HandleRecvChunk 驱动 CopyToDeviceStream 来提供数据，即 CopyToDeviceCallback）。这些名字与直觉读者的预期相反。


PJRT transfer manager	`xla::TpuHostTransferManager` (`tpu_pjrt_client.cc`), ctor @ `0xf8150c0`, vtable `off_2177B988`
PJRT execute 注册点	`xla::TpuHostTransferManager::SetExecuteEvent(AsyncValueRef<tpu::TpuEvent>)` @ `0xf813760`, called from `TpuExecutableLoadState::ExecuteLaunchRaw` @ `0xf8109a0`
PJRT send (dev→host) chunk	`xla::TpuHostTransferManager::HandleSendChunk(uint32_t channel, tsl::Future<PjRtChunk>)` @ `0xf815720`
PJRT recv (host→dev) chunk	`xla::TpuHostTransferManager::HandleRecvChunk(uint32_t channel, unique_ptr<CopyToDeviceStream>)` @ `0xf815a20`
PJRT 回调注册	`xla::(anon)::SetUpHostCallbacksForDevice(Span<SendCallback>, Span<RecvCallback>, flat_hash_map<long,PjRtTransferMetadata>, flat_hash_map<long,vector<long>>, TpuClient, TpuDevice, bool, TpuHostTransferManager*)` (`tpu_pjrt_client.cc`)
PJRT channel 键	`uint32_t` channel id → `absl::flat_hash_map` (send map slots @ `this+8/16/24/32`, recv map slots @ `this+40/48/56/64`)
PJRT 排空线程	`TpuHostTransferManagerSendThread` / `TpuHostTransferManagerRecvThread` (two detached `concurrent::LoopExecutor`, @ `this+192` / `this+200`)
遗留 transfer manager	`tensorflow::TpuHostTransferManagerBase` (`tpu_host_transfer_manager.cc`); impl `TpuHostTransferManagerImpl` ctor @ `0xeab4080`
遗留 init / rendezvous	`TpuHostTransferManagerBase::Initialize(Span<HostTransferProto>, rendezvous_key_base, …, RendezvousInterface*, TpuTopology)` @ `0xeab6780`
遗留 device→host 分发	`TpuHostTransferManagerBase::HandleHostCommand(HostCommandParams, TpuStackBases)` @ `0xeab5b20` → `HandleSendSF(int,long)` @ `0xeab5520` / `HandleRecvSF(int,long)` @ `0xeab5780`
遗留 channel 键	`int` channel id → `std::map<int, HostSendRecvInfo>` (red-black tree @ `this+24`)
遗留 StreamExecutor shim	`tpu::TpuStream::DoHostCallbackWithStatus` @ `0xe998fa0` → `HostCallbackTrampoline` @ `0xe999660` (ExecutorApiFn `+440` enqueue, `+368` status return)
Driver send/recv 叶子节点	`deepsea::executor::DeepseaStream::EnqueueOnDeviceSendRecvLocal` / `EnqueueOnDeviceSend` / `EnqueueOnDeviceRecv` (TPU driver core, opaque)
证据等级	Reimplementation-grade / byte-confirmed against IDA decompile (both managers traced; channel keying, registration, rendezvous, and fatal-on-miss confirmed)

1. 两个 Transfer Manager，一个 Channel 模型

libtpu 像处理 infeed/outfeed 一样，也携带了两套主机回调机制。经字节确认的拆分如下：

方面	PJRT 路径	遗留 StreamExecutor / TF 路径
类	`xla::TpuHostTransferManager`	`tensorflow::TpuHostTransferManagerBase` (+ `…Impl`)
生命周期	每次 execution launch 一个	每次 `LocalClient` executable load 一个
构造参数	`{TpuClient*, TpuCoreLocation}` (`0xf8150c0`)	`{int ordinal, xla::Backend, StreamExecutor}` (`0xeab4080`)
Channel 键	`uint32_t` → two `absl::flat_hash_map` (send / recv)	`int` → `std::map<int, HostSendRecvInfo>`
注册	execute 路径上的 `SetExecuteEvent` (`0xf813760`)；回调由 `SetUpHostCallbacksForDevice` 安装	`Initialize(rendezvous_key_base, RendezvousInterface*, …)` (`0xeab6780`)
设备侧触发	TPU completion fulfils `tsl::Future<PjRtChunk>` (send) / drives `CopyToDeviceStream` (recv)	TPU driver 使用打包命令字触发 `HandleHostCommand`
匹配	对 channel id 做 hash-map probe	`RendezvousInterface::Send` / `RecvAsync(parsed_key, …)`
排空	两个 detached 线程 + `concurrent::LoopExecutor`	在 driver-callback 线程 + TF rendezvous 上 inline
源码根	`learning/45eac/research/pjrt/tpu_pjrt_client.cc`	`learning/45eac/google/xla/tpu_host_transfer_manager.cc`
Driver 叶子节点	`tpu::System` events → `DeepseaStream::EnqueueOnDevice*`	`tpu::TpuStream` → `ExecutorApiFn+440` → `DeepseaStream::EnqueueOnDevice*`

GOTCHA — 它们不是分层关系。 PJRT TpuHostTransferManager 从不构造 TpuHostTransferManagerBase，也从不接触 TF RendezvousInterface；它在两个 flat_hash_map 中自行做 channel 匹配，并 fulfil tsl::AsyncValue。遗留 TpuHostTransferManagerBase 从不直接引用 TpuClient / tpu::System；它经由 StreamExecutor TpuStream 和 TF rendezvous。假设 PJRT 路径通过遗留 manager 转发的重新实现者会出错。二者只在 TPU driver core 的 on-device send/recv（DeepseaStream::EnqueueOnDevice*）内部汇合，但抵达那里的是不同的最后一跳。置信度：CONFIRMED。

还有第三个 TFRT 风格变体 tensorflow::tfrt_tpu::HostTransferManager（TransferDeviceToHost(int, MutableBorrowingLiteral*, function<void(Status)>) @ 0xe70bae0，TransferHostToDevice(int, BorrowingLiteral const&, …) @ 0xe70c5c0）—— 一个以 int channel 为键、literal 形态的薄包装，会转发进同一套机制。它是 chunk-API PJRT 路径的 literal-API 类比，不在主 execute 热路径上；本页记录的是基于 chunk 的 PJRT manager 和遗留 rendezvous manager。

2. Channel-ID 键与 Rendezvous 契约

2.1 什么命名一个回调

主机回调不携带设备地址或队列索引。它由 HLO Send/Recv op 携带的整数 channel id 命名（即 XLA 写入 op 的同一个 channel_id 字段）。编译器围绕主机计算发出 Send/Recv 对；设备侧（op）和主机侧（已注册回调）都命名同一个 channel id，运行时据此匹配。

PJRT 以 uint32_t 为键。每个 TpuHostTransferManager 都持有两个内联布局在对象中的 absl::flat_hash_map：

Map	方向	HLO op	回调角色	对象中的槽位
Send map	device → host	`Send`	`CopyFromDeviceCallback`（主机接收 `PjRtChunk`）	size `this+8`, seed `this+16`, ctrl `this+24`, slots `this+32`
Recv map	host → device	`Recv`	`CopyToDeviceCallback`（主机通过 `CopyToDeviceStream` 提供数据）	size `this+40`, seed `this+48`, ctrl `this+56`, slots `this+64`

QUIRK — "Send" 是 device→host，"Recv" 是 host→device。 命名来自设备程序的视角：程序的 Send op 把数据 off-chip 发送给主机（所以主机侧回调是 CopyFromDeviceCallback，chunk 流向主机）；程序的 Recv op 从主机接收数据（所以主机侧回调是 CopyToDeviceCallback，主机把 chunk 下推到 CopyToDeviceStream）。如果重新实现者把 HandleSendChunk 接到 host-input 路径、把 HandleRecvChunk 接到 host-output 路径，数据方向会完全反过来。置信度：CONFIRMED — 两个回调 map 及其 fatal-miss 消息明确命名了方向（§2.3）。

遗留路径以 int 为键。TpuHostTransferManagerBase 持有单个 std::map<int, HostSendRecvInfo>（一棵根在 this+24 的红黑树），由 Initialize 期间的 AddHostTransfer 从 HostTransferProto 填充。HostSendRecvInfo 携带 channel id、to_host 方向标志（偏移 +40）、sync-flag 值（+166）以及解析后的 rendezvous key。

2.2 PJRT 注册：`SetUpHostCallbacksForDevice` → `SetExecuteEvent`

对于一次 PJRT execute，用户传入 ExecuteOptions::send_callbacks / recv_callbacks（SendCallback / RecvCallback 的 span）。xla::(anon)::SetUpHostCallbacksForDevice 是注册入口 —— 其签名就是契约：

// xla::(anon)::SetUpHostCallbacksForDevice   (tpu_pjrt_client.cc)
// builds the per-launch TpuHostTransferManager's two channel maps
SetUpHostCallbacksForDevice(
    Span<const SendCallback>        send_callbacks,   // one per dev->host channel
    Span<const RecvCallback>        recv_callbacks,   // one per host->dev channel
    const flat_hash_map<long, PjRtTransferMetadata>&  metadata_by_channel,  // device/host shapes
    const flat_hash_map<long, vector<long>>&          channel_groups,
    TpuClient*, TpuDevice*, bool use_major_to_minor_data_layout_for_callbacks,
    TpuHostTransferManager* mgr);   // the manager whose maps get populated
```text

每个 `SendCallback` 都会成为 send map 中以其 channel id 为键的条目；存储的 closure 会以接收到的 `PjRtChunk` 调用（`__call_func<…SetUpHostCallbacksForDevice…::$_0>` thunk @ `0xf81f900` 是 send-callback 主体 —— 它在调用用户回调前将设备 chunk delinearize 为主机布局）。每个 `RecvCallback` 都会成为 recv map 中的条目（`…::$_1` thunk 处理 `CopyToDeviceStream`）。

manager 由 `TpuExecutableLoadState::ExecuteLaunchRaw`（@ `0xf8109a0`）绑定到 launch，该函数会在 `tpu::System::Execute` 前立即调用 `TpuHostTransferManager::SetExecuteEvent`：

```c
// xla::TpuHostTransferManager::SetExecuteEvent   sub_F813760
// a1 = TpuHostTransferManager, a2 = &AsyncValueRef<tpu::TpuEvent>  (the launch's completion event)
core_loc   = a1+96;                                       // the {TpuCoreLocation} copied at ctor
shm        = TpuCoreLocation::LocalSharedMemory(core_loc, 0);
logger     = client.system().pending_event_loggers();     // *(*(a1+88)+648)+160
if (logger && (lg = logger->get(shm.index_on_host())))    // attach a "done" marker to the event
    lg->vtable+56(done_async_value@a1+176, &execute_event);   // chain the manager's done AV
// then enqueue this manager as a waiter on the execute TpuEvent:
issuer = lock(weak_ptr @ a1+80);                          // throws bad_weak_ptr if expired
if (execute_event.IsConcrete())                           // already done?
    a1->vtable+40(a1, status_from(execute_event));        // fire MaybeCallDoneCallback now
else
    AsyncValue::EnqueueWaiterListNode(execute_event, node, …);  // fire when execution completes

因此 SetExecuteEvent 会使 transfer manager 成为 launch 的 tpu::TpuEvent 的一个 waiter：执行结束（或出错）时，manager 的 done 路径运行（MaybeCallDoneCallback @ 0xf815500），并在所有未完成 chunk 传输排空后销毁 manager。这就是每次 launch 的 transfer manager 生命周期如何绑定到 execute event，而不是绑定到调用方。

2.3 PJRT chunk 分发与 fatal-on-miss

当设备遇到 Send 时，TPU completion fulfil 一个 tsl::Future<PjRtChunk>，send-drain 线程调用 HandleSendChunk(channel, future)。当设备遇到 Recv 时，recv-drain 线程调用 HandleRecvChunk(channel, CopyToDeviceStream)。二者都会在 channel id 上做 hash-map probe，miss 时都会 CHECK-fail：

// xla::TpuHostTransferManager::HandleSendChunk   sub_F815720   (dev->host)
// a1 = this, a2 = uint32 channel, a3 = &tsl::Future<PjRtChunk>
mutex.lock(this+184); ++this->outstanding_count[this+168]; mutex.unlock(this+184);
slot = flat_hash_map_find(send_map @ this+8, /*key=*/a2);   // CRC32 hash + SIMD ctrl probe
if (!slot)                                                  // tpu_pjrt_client.cc:4712
    LOG(FATAL) << "No CopyFromDeviceCallback registered for channel " << a2;
issuer = lock(weak_ptr @ this+72);                          // bad_weak_ptr if expired
RET_CHECK(future.IsValid());                                // future.h:420
future.AndThen([slot.callback](StatusOr<PjRtChunk> chunk){ /* run user send callback */ });

// xla::TpuHostTransferManager::HandleRecvChunk   sub_F815A20   (host->dev)
// a1 = this, a2 = uint32 channel, a3 = unique_ptr<CopyToDeviceStream>
slot = flat_hash_map_find(recv_map @ this+40, /*key=*/a2);  // separate map, same probe
if (!slot)                                                  // tpu_pjrt_client.cc:4732
    LOG(FATAL) << "No CopyToDeviceCallback registered for channel " << a2;
mutex.lock(this+184); ++this->outstanding_count[this+168]; mutex.unlock(this+184);
wrapper = operator new(0x2C8);                              // wraps the CopyToDeviceStream
wrapper->vtable = off_2177CF00; wrapper->stream = a3; …
post wrapper to recv LoopExecutor @ this+200;              // invoke recv callback off-thread
```text

> **GOTCHA — 未注册 channel 是致命崩溃，不是 no-op。** 如果 channel id 不存在于对应 map，`HandleSendChunk`（`:4712`）和 `HandleRecvChunk`（`:4732`）都会调用 `LogMessageFatal`。重新实现者必须在 launch 前注册编译后程序会使用的*每个* channel；没有回调的 `Send`/`Recv` op 会中止进程。遗留路径也是一样 —— `HandleSendSF` 会以 `"Channel id not found in transfer map."`（`tpu_host_transfer_manager.cc:110`）中止。置信度：CONFIRMED。

`this+168` 处的 `outstanding_count`（由 `this+184` 处的 mutex 保护）跟踪正在进行的 chunk 传输；manager 的 done 回调（`MaybeCallDoneCallback`）只有在它排空到零后才完成 launch，所以慢速主机回调会保持 launch 的 completion event 处于未完成状态。

### 2.4 遗留 rendezvous：`Initialize` → `HandleHostCommand`

遗留 manager 从编译后程序的 `HostTransferProto` 和 `rendezvous_key_base` 字符串设置：

```c
// tensorflow::TpuHostTransferManagerBase::Initialize   sub_EAB6780
// (Span<HostTransferProto> transfers, string_view rendezvous_key_base,
//  string_view device_type, RendezvousInterface* rendezvous, TpuTopology)
collect transfers into InlinedVector<HostTransferProto*,4>;     // 120 B per proto
log "TpuHostTransferManagerBase::Initialize rendezvous_key_base=" << rendezvous_key_base;
for (proto : transfers)
    AddHostTransfer(proto, &transfer_map /*flat_hash_map<string, HostTransfer*>*/, topology);
    // each AddHostTransfer parses the rendezvous key:
    //   "host_compute_rendezvous:" + key, SimpleAtoi(*rendezvous_key, &channel_int),
    //   inserts {channel_int -> HostSendRecvInfo} into the std::map @ this+24
store rendezvous_ @ this+5;                                      // the TF RendezvousInterface*

运行时 TPU driver 会用打包的 32 位命令字触发 HandleHostCommand：

// tensorflow::TpuHostTransferManagerBase::HandleHostCommand   sub_EAB5B20
// a2 = &command_word, a3 = &TpuStackBases  (sync-flag pointer)
cmd  = *a2;
type = cmd >> 24;            // high byte: 1 = Send, 2 = Recv
chan = cmd & 0xFFFFFF;       // low 24 bits: channel id
if (type == 1) { HandleSendSF(chan, *a3); return true; }   // dev->host
if (type == 2) { HandleRecvSF(chan, *a3); return true; }   // host->dev
return false;                                              // unknown command -> not handled
```text

`HandleSendSF`（`0xeab5520`）按 channel 遍历 `std::map<int, HostSendRecvInfo>`，CHECK 该 transfer 为 `to_host`（`:119`），将 sync flag 记录进 `HostSendRecvInfo`，并通过 TF rendezvous 匹配主机侧：

```c
// tensorflow::TpuHostTransferManagerBase::HandleSendSF   sub_EAB5520
info = transfer_map.find(channel);                          // rbtree walk @ this+24
if (!info) LOG(FATAL) << "Channel id not found in transfer map.";   // :110
CHECK(info->to_host);                                       // :119
info->sync_flag_slot->sflag = info->sflag_value;            // *(v4[78]+28) = *(v4+166)
info->sync_flag_slot->ptr   = sync_flag_ptr;                // *(v4[78]+32) = a3
++this->pending_count;                                      // InterlockedIncrement
TF_CHECK_OK(rendezvous_->Send(info->parsed_key, args, val, /*is_dead=*/false));   // :137

HandleRecvSF 是 rendezvous_->RecvAsync(parsed_key, …) 上的镜像，会完成到 RendezvousDone（0xeab59e0）。rendezvous key 字符串是标准的 XLA host-compute key：host_compute_channel_{0}_args / host_compute_channel_{0}_retvals，并以 host_compute_rendezvous: 为前缀。

3. StreamExecutor 主机回调 Shim

在遗留 TpuHostTransferManagerBase 之下（并且任何 StreamExecutor consumer 都可抵达）是 Stream::DoHostCallbackWithStatus(absl::AnyInvocable<absl::Status()>) 原语 —— 最低层的“在主机上运行此 closure”钩子。它有两个后端。

Host (CPU) —— inline。 stream_executor::host::HostStream::DoHostCallbackWithStatus @ 0xfe6efe0 会立即在调用线程上运行 closure（return (*callback)();）。链接的 HostStream 是同步变体，所以 BlockHostUntilDone（0xfe6f000）直接返回 true。这是 host-memory-staging / trivial-host-op 路径，不是主机回调 rendezvous。

TPU —— C-shim trampoline。 tensorflow::tpu::TpuStream::DoHostCallbackWithStatus @ 0xe998fa0 是遗留 outside-compilation 的实现：

// tensorflow::tpu::TpuStream::DoHostCallbackWithStatus   sub_E998FA0
holder = operator new(32, 16);                              // closure holder
move AnyInvocable(callback) into holder+0x10;              // 16-byte trivial state + manager/invoker
// enqueue on the TPU command stream; driver fires it once preceding work completes:
ExecutorApiFn()+440( SE_Stream@this+0x88, stream_handle@this+0x80,
                     &TpuStream::HostCallbackTrampoline, holder );
// on failure: MakeErrorImpl<13>("Failed to  host callback.")   (tpu_stream.h:177)

// tensorflow::tpu::TpuStream::HostCallbackTrampoline   sub_E999660   (the C-side fire)
status = (*holder->closure)();                              // run user closure -> absl::Status
ExecutorApiFn()+368( status.code(), StatusMessageAsCStr(status) );  // report into SE_Status
destroy holder->closure; operator delete(holder, 32, 16);
```text

> **NOTE — ExecutorApiFn 槽位就是契约。** 遗留主机回调路径中 libtpu 可见的唯一旋钮是 `TfTpu_ExecutorApiFn` 表偏移：**+440** = 在 TPU stream 上 enqueue 主机回调，**+368** = 将生成的 status 回报到 driver 的 `SE_Status`。这些是 TPU-driver C-ABI 的一部分（单独任务；不是完整表）。当 TPU stream 抵达 enqueued callback 时，driver 暂停 stream，在主机线程上运行 trampoline，传回 `absl::Status`，然后继续。重新实现遗留路径时，需要复现 32 字节 closure-holder 生命周期和这两个槽位调用。置信度：槽位编号和 closure 生命周期为 CONFIRMED；driver 内部 stream-pause/resume 对 `libtpu.so` 不透明（精确 driver 机制为 LOW）。

---

## 4. On-Device Send/Recv（Driver 深度）

两个 manager 最终都驱动 TPU driver core 内相同的 on-device send/recv。channel 的设备侧是一个 **sync flag**（`tpu::TpuSyncFlagOnChip`）：`Send` op 写入数据并升起主机等待的 sync flag；`Recv` op 阻塞在 sync flag 上，直到主机交付数据后升起它。driver 叶子节点如下：

| Driver 符号 | 角色 |
|---|---|
| `deepsea::executor::DeepseaStream::EnqueueOnDeviceSendRecvLocal(DeviceAddressBase, DeviceAddressBase)` | local same-host send↔recv（快速 on-device 路径） |
| `deepsea::executor::DeepseaExecutor::EnqueueOnDeviceSendLocal(DeviceAddressBase, DeviceAddressBase, TpuSyncFlagOnChip, AnyInvocable<void(Status)>)` | device→host send，带 completion callback |
| `deepsea::executor::DeepseaExecutor::EnqueueOnDeviceRecv(DeviceAddressBase, function<void(TpuSyncFlagOnChip)>, AnyInvocable<void(Status)>)` | host→device recv |
| `OnDeviceSendRecvLocalRequest::ExecuteImpl(DeepseaStream*)` | driver 分发的排队 request 对象 |

PJRT 路径通过 `tpu::System` events 抵达这些节点（manager fulfil driver 消费的 `tsl::AsyncValue`）；遗留路径通过 `TpuStream` + `ExecutorApiFn` 抵达。重新实现者从这一层只需要知道：(a) 设备侧是 sync-flag wait/raise，不是阻塞 memcpy；(b) "local" send/recv（`EnqueueOnDeviceSendRecvLocal`）是同一主机上的优化快捷路径，当两端在同一主机上时会绕过主机回调（`"EnqueueOnDeviceSend using same host fast path."`）；(c) 实际的 silicon sync-flag 写入和 stream pause/resume 在 TPU driver core 内部，对这个二进制不透明。置信度：符号集合和 sync-flag 模型为 HIGH；silicon 机制为 LOW。

> **QUIRK — host_compute 可以被禁用，从而把 Send/Recv 变成 host-to-host copy。** 标志 `xla_disable_automatic_host_compute_offload`（字符串：*"Allow host-to-host copy when automatic host compute offload is disabled…"*) 会改变编译器是否实际 offload 主机计算。关闭 automatic offload 时，运行时可能通过 host-to-host copy 而不是 device-mediated rendezvous 来满足一个 channel。只建模 device-rendezvous 路径的重新实现者会错误处理 disabled-offload 配置。置信度：HIGH（字符串 + flag 已确认；精确 host-to-host 路径未追踪）。

---

## 5. 重新实现说明

- **选择正确的 manager。** 传入 `send_callbacks` / `recv_callbacks` 的 PJRT consumer（JAX / PyTorch-XLA）使用 `xla::TpuHostTransferManager`，由 `SetExecuteEvent` 按 launch 注册。只有 `LocalClient` / TF-TPU op kernel 使用带 TF `RendezvousInterface` 的 `tensorflow::TpuHostTransferManagerBase`。实现其中一个并不等于实现另一个。
- **键是 channel id，而且是必需的。** 将回调表建模为 `map<channel_id, callback>`；PJRT 上有两个表（send / recv），以 `uint32_t` 为键；遗留路径上有一个表，以 `int` 为键。launch 前注册每个 channel —— 未匹配的设备命令是 `LOG(FATAL)`，不是被丢弃的消息。
- **注意方向反转。** `Send`（HLO）= device→host = `HandleSendChunk` = `CopyFromDeviceCallback`（主机接收 `PjRtChunk`）。`Recv`（HLO）= host→device = `HandleRecvChunk` = `CopyToDeviceCallback`（主机通过 `CopyToDeviceStream` 推送）。按 HLO op 语义接线，不要按英文词面接线。
- **注册将 manager 绑定到 execute event。** `SetExecuteEvent` 使 manager 成为 launch 的 `tpu::TpuEvent` 的 waiter；只有所有未完成 chunk 传输都排空后（`this+168` 的 `outstanding_count`，由 `this+184` 的 mutex 保护），manager 的 done 回调才会完成 launch。慢速主机回调会延迟 launch 的 completion future。
- **PJRT 上有两个排空线程。** PJRT manager 会生成 `TpuHostTransferManagerSendThread` 和 `TpuHostTransferManagerRecvThread`（位于 `this+192` / `this+200` 的 detached `concurrent::LoopExecutor`）；chunk 回调在这些线程上运行，脱离 device-completion 线程。遗留路径通过 rendezvous 在 driver-callback 线程上运行回调。
- **遗留命令字是打包的。** `HandleHostCommand` 解码 `{cmd>>24 = Send|Recv, cmd&0xFFFFFF = channel}`。复现这个 8/24 位拆分；高字节是方向，低 24 位是 channel。
- **这是按 channel 键控的 rendezvous，不是 FIFO。** 对比 [Infeed / Outfeed 队列](infeed-outfeed.md)：infeed/outfeed 命名一个 `{TpuCoreLocation, queue_index}`，设备按程序顺序消费条目，没有逐次传输身份；主机回调命名 channel id，运行时将设备 `Send`/`Recv` 与为该确切 channel 注册的主机回调匹配。`xla.host_compute` / outside-compilation 和跨主机 `send`/`recv` 使用主机回调；流式批量输入/输出使用 infeed/outfeed。

---

## 相关组件

| 名称 | 关系 |
|---|---|
| `xla::TpuHostTransferManager` | PJRT per-launch channel-keyed callback manager（本页主题） |
| `tensorflow::TpuHostTransferManagerBase` | 遗留 TF-rendezvous channel-keyed manager（本页主题） |
| `tensorflow::tfrt_tpu::HostTransferManager` | TFRT literal-API 变体（`TransferHostToDevice`/`TransferDeviceToHost(int channel, …)`） |
| `TpuExecutableLoadState::ExecuteLaunchRaw` | 调用 `SetExecuteEvent` 来注册 manager 的 execute 叶子节点 |
| `tpu::TpuStream::DoHostCallbackWithStatus` | 遗留路径可抵达的 StreamExecutor 主机回调 shim |
| `deepsea::executor::DeepseaStream::EnqueueOnDeviceSendRecvLocal` | 两条路径最终落到的 TPU-driver on-device send/recv 叶子节点 |

主机回调 ​

摘要 ​

1. 两个 Transfer Manager，一个 Channel 模型 ​

2. Channel-ID 键与 Rendezvous 契约 ​

2.1 什么命名一个回调 ​

2.2 PJRT 注册：SetUpHostCallbacksForDevice → SetExecuteEvent ​

2.3 PJRT chunk 分发与 fatal-on-miss ​

3. StreamExecutor 主机回调 Shim ​