Megascale — 章节地图

地址适用于 libtpu-0.0.40-cp314 wheel 中的 libtpu.so。其他版本会有所不同。二进制： extracted/libtpu-0.0.40-cp314-cp314-manylinux_2_31_x86_64/libtpu/libtpu.so（build-id 89edbbe81c5b328a958fe628a9f2207d；.text VMA == 文件偏移）。下面所有符号都存在于完整符号二进制中；已将 demangled 名称和地址与 IDA 反编译结果交叉核对。

摘要

Megascale 是 TPU runtime 的数据中心网络（DCN）横向扩展层。ICI 是 pod 内光学环面网络，把一个 slice 的芯片连成单个 SerDes 连接的 (X,Y,Z) mesh；Megascale 则把许多这样的 slice（每个 slice 都是自己的 ICI 岛）联合成一个跨越数据中心机架和行的逻辑 TPU 作业，并通过以太网上的 gRPC 通信。ICI 在一个 slice 内以链路速度在芯片之间移动操作数；Megascale 在 slice 之间的主机之间移动 collective 的跨 slice 部分（编译器在 xla::megascale::compiler::CrossSliceRewrites 中发出的 AllReduceDcnFusionData / LowerCollectivePermuteFullDCN 路径），并拥有让这些主机能够互相发现的控制平面。这种分层在二进制自己的词汇中也可见：跨 slice rewriter 将其原语命名为 ICIDCN… 和 …FullDCN，把 “ICI” 和 “DCN” 视为同一个 collective 的两个传输层级。

本章节的所有内容都位于 C++ namespace xla::megascale::runtime 下，并由一个门面对象 CommunicationBackend 承载（对象大小 0x370，由 Create() 构建）。每个 TPU 主机运行一个 backend；该 backend 拥有一个 gRPC server，即 6 方法的 MegaScaleTransport service，绑定到 MEGASCALE_PORT，并且在每个作业中正好一个被选举出的进程上，还拥有一个 TopologyCoordinator、一张 BarrierCoordinator 映射，以及一个 ErrorReporter。在 backend 之下，主机本地的 tpunetd daemon（namespace superpod::tpunetd_client）会在 Megascale 的跨 slice rendezvous 运行之前启动 slice 内部的 ICI fabric；二者按单向顺序排列，即 tpunetd → Megascale。

本页是一张地图。它帮助读者定位 ICI↔DCN 边界、主机侧 daemon 模型，以及四阶段作业生命周期，然后把每个字节级推导交给已经写好的子章节。它不重复这些子章节的内部细节。

对于整个章节，重新实现契约是：

两个传输层级及其分界点 — ICI（slice 内 torus，由 tpunetd 拥有）vs DCN（跨 slice 以太网/gRPC，由 Megascale 拥有），以及从一个层级跨到另一个层级的唯一数据（TpuTopologyArgsProto）。
主机侧 daemon 模型 — 每个主机一个 CommunicationBackend，每个作业一个被选举出的 coordinator，MegaScaleTransport gRPC 表面，以及其下方的主机本地 tpunetd。
作业生命周期 — rendezvous → fleet metadata 分发 → 跨主机 barrier → 稳态 heartbeat，并以错误聚合作为与四个阶段正交运行的失败主干。


C++ namespace	`xla::megascale::runtime`
每主机门面	`CommunicationBackend`（`0x370` 字节；`Create()` @ `0x1ccafe60`）
Bootstrap 入口	`CommunicationBackend::DiscoverTopologyAndAddressBindings(int, tpu::TpuTopologyArgsProto, int, int)` @ `0x1ccacb80`
gRPC service	`MegaScaleTransport` — 6 个 unary RPC，前缀 `/xla.megascale.runtime.MegaScaleTransport/`
Coordinator（1/job）	`TopologyCoordinator`（`0x108` 字节）+ `BarrierCoordinator` map + `ErrorReporter`
下方（slice 内）	通过 `superpod::tpunetd_client` 访问的 `tpunetd` daemon（ICI fabric，UDS）
选举开关	`MEGASCALE_COORDINATOR_ADDRESS` / `--megascale_coordinator_address`
源码根目录	`platforms/xla/megascale/runtime/communication/`

ICI↔DCN 边界

这是本章节围绕的核心区别，二进制也把它划分得很清楚。

两套坐标系统，在一个点相接

一个 TPU 作业是两层网络彼此堆叠形成的层级结构：

text

   ┌──────────────────────────────────────────────────────────────┐
   │  DCN 层级 — Megascale (xla.megascale.runtime)                  │
   │  不同 slice 的主机之间使用 Ethernet + gRPC。                  │
   │  用 (slice_id, host_id) + 每 slice device id 命名芯片。        │
   │  拥有：CommunicationBackend, MegaScaleTransport, coordinators. │
   └──────────────────────────────┬─────────────────────────────────┘
                                  │  TpuTopologyArgsProto（唯一
                                  │  跨越边界的数据）
   ┌──────────────────────────────┴─────────────────────────────────┐
   │  ICI 层级 — 每个 slice 一个光学 torus (superpod.routing.*)     │
   │  一个 slice 内芯片之间的 SerDes 链路。                         │
   │  用 ToroidalTopology 中的 ChipCoordinate (X,Y,Z) 命名芯片。    │
   │  拥有：tpunetd link bring-up, routing tables, GTC sync.        │
   └──────────────────────────────────────────────────────────────────┘
```text

这两层会刻意保持各自的寻址方案相互分离。ICI 层用 [topology discovery](../ici/topology-discovery.md) 的 square-seed polarity inference 产生的笛卡尔 `ChipCoordinate` 寻址芯片；DCN 层用 `(slice_id, host_id)` 寻址一个 *host*，并且绝不重新推导芯片坐标。从 ICI 向上跨到 DCN 的唯一数据是 **每 slice shape**，它序列化为 `tpu::TpuTopologyArgsProto`，并逐字携带在每个 `GetMultiSliceTopologyRequest`（字段 2）和每个 `SliceInfo`（字段 3）中。DCN 层把该 shape 视为不透明且已验证的 blob；coordinator 只会对它运行 `proto2::util::MessageDifferencer::Compare` 来确认所有 slice 一致，而不会解码芯片几何结构。（验证调用位于 `TopologyCoordinator::ProcessRequest` @ `0x1cf524c0`。）

> **注意 —** 反编译也在 *collective* 层面把双层模型具体化。编译器 pass `CrossSliceRewrites::LowerAllReduce`（@ `0x111d6f40`）接收一个 `optional<AllReduceDcnFusionData>`，而 `ConstructAllGatherICIDCNextReceiverConstant`（@ `0x14b96340`）在同一个符号名中命名了两个层级：跨 slice 的 all-reduce 被融合为一个 *slice 内 ICI 阶段* 和一个 *跨 slice DCN 阶段*。Megascale 拥有 DCN 阶段；ICI 拥有其余部分。

### 为什么要集中化 DCN 控制平面

ICI discovery 是去中心化的；slice 中的每个芯片都会通过 BFS 遍历其邻居的链路签名来推断自己的坐标（参见 [Topology Discovery](../ici/topology-discovery.md)）。DCN rendezvous 则相反：**严格集中化**，每个作业一个 coordinator 进程。原因是可达性。ICI 链路是物理接线且自描述的；芯片可以在没有带外通道的情况下发现自己的 torus 邻居。DCN peer 是散布在数据中心中的以太网主机，彼此的地址没有先验知识。Megascale 用经典 parameter server 的方式解决这个问题：一个进程绑定众所周知的 endpoint（`MEGASCALE_COORDINATOR_ADDRESS`），其他所有进程连接到它，然后 coordinator 组装并广播单个集群范围的地址表。这里没有 peer discovery，也没有投票；选举纯粹是“谁的 `MEGASCALE_COORDINATOR_ADDRESS` 解析到本地接口”。

---

## 主机侧 Daemon 模型

作业中的每个 TPU 主机都运行相同的两个软件层；只有 coordinator 角色不同。

### CommunicationBackend — 每个主机一个

`CommunicationBackend`（`0x370` 字节，由 `Create()` @ `0x1ccafe60` 通过 `operator new(880, 16)` 分配）是所有 DCN 活动的每进程门面。它拥有：

- `MegaScaleTransport` gRPC server，绑定到 `MEGASCALE_PORT`，并由 `InitializeTransportLayerInternal()` @ `0x1ccaeb40` 通过 `GrpcTransport`（默认）连线；
- 每 `(slice,host)` 的 **地址表**（成员位于 `+0x170`，在 `+0x1a0`/`+0x1a8` 两个指针槽下方），供稳态 `Send` RPC 索引；
- 位于 `+0x1a0` 的 `TopologyCoordinator*` — **除 coordinator 外的每个进程上都为 null**；
- 位于 `+0x1b0` 的 `flat_hash_map<string, unique_ptr<BarrierCoordinator>>`，以 `barrier_id` 为 key，并由 `+0xe0` 的 `TracedMutex` 保护；
- 位于 `+0x1a8` 的 `ErrorReporter*`；
- 由 `StartHeartBeat()` @ `0x1ccade60` 启动的 `HeartBeat` scheduler。

这种非对称性就是整个设计：worker 的 `+0x1a0` 槽位保持 null，因此当一个误发的 Topology RPC 到达它时，`OnTopologyRequestReceived` @ `0x1ccac380` 会用 `"TopologyCoordinator not initialized."`（INTERNAL，category 13）拒绝它。只有被选举出的进程才会运行 `InitializeCoordinator(num_slices)` @ `0x1ccad600`，它会 `new` 出 `TopologyCoordinator` 并构造每进程的 `ErrorReporter`。

### MegaScaleTransport — 6 方法 gRPC 表面

每个跨主机控制消息（bootstrap 和稳态）都通过一个 unary gRPC service 传输。反编译通过 `WithCallbackMethod_*` 模板链确认了 server 侧嵌套（Send → GetMultiSliceTopology → Barrier → ReportError → TriggerError → …，可见于 `0x1ce80200`+ 的符号）：

| RPC | Request → Response | 阶段 | 所属页面 |
|---|---|---|---|
| `GetMultiSliceTopology` | `GetMultiSliceTopologyRequest` → `…Response` | bootstrap | [Bootstrap](bootstrap/overview.md) |
| `Barrier` | `BarrierRequest` → `BarrierResponse` | barrier | [跨主机 Barrier](cross-host-barrier.md) |
| `ReportError` | `ReportErrorRequest` → `…Response` | 失败 | [错误聚合器](error-aggregator.md) |
| `TriggerError` | `TriggerErrorRequest` → `…Response` | 失败 | [错误聚合器](error-aggregator.md) |
| `Send` | `SendRequest` → `TransferBufferResponse` | 稳态 | —（DCN collective transfer） |
| `SendHeartBeat` | `HeartBeatRequest` → `HeartBeatResponse` | 稳态 | —（liveness） |

该 service 的 wire protocol，包括 message taxonomy、transport selection（`MEGASCALE_TRANSPORT_TYPE` = `grpc` default vs `chaotic_good_legacy`）、authentication，以及 `GrpcTransport` server wiring，归 [tpunetd Protocol](tpunetd-protocol.md) 和 [Bootstrap › Worker Registration](bootstrap/worker-registration.md) 所有。

### tpunetd — 下方的 slice 内 daemon

`tpunetd` 是主机本地 daemon，不属于 `libtpu.so` 自身；`libtpu` 通过 Unix domain socket 上的 `superpod::tpunetd_client` 与它通信。它拥有 *一个 slice 内部的一切*：ICI link 配置、routing-table 生成、global-time-counter（GTC）同步、每主机 `SetChipCoordinates`，以及 slice 内 `SessionMaster::BroadcastBarrier` rendezvous。Megascale 严格运行在它 *之上*，并消费它的输出：每个 worker 放入其 `GetMultiSliceTopologyRequest` 的 `HostNetworkAddress` 和 `TpuTopologyArgsProto` 都派生自 tpunetd 的芯片坐标状态，因此 tpunetd 必须在 `DiscoverTopologyAndAddressBindings` 能够运行之前完成。该依赖是单向的；tpunetd 从不阻塞在 `MegaScaleTransport` 上，并且当 `enable_megascale_topology=false`（单 slice 模式）时，`MegaScaleTransport` server 永远不会启动，但 tpunetd 仍然会运行。完整交接见 [Bootstrap › ICI Handoff](bootstrap/ici-handoff.md) 和 [Bootstrap › tpunetd Relationship](bootstrap/tpunetd-relationship.md)。

> **易错点 —** PJRT distributed `CoordinationService`（package `xla.coordination`）是共驻于 `libtpu` 中的 *独立* 控制平面，不属于 Megascale。它为 JAX/TF 协调 Python 层状态（shard assignment、run id）；`MegaScaleTransport` 协调 TPU 内部芯片 fabric 状态。它们使用不同的 proto namespace、不同的端口，并且关键的是，**绝不共享 barrier ID**。重新实现者如果混淆二者，就会在错误的 rendezvous 上等待而死锁。`xla.coordination.BarrierRequest` 和 `xla.megascale.runtime.BarrierRequest` 是不同的 protobuf 类型。

---

## 作业生命周期

一次完整的 Megascale 作业启动包含四个有序阶段，并以错误聚合作为第五条正交主干，任一阶段都可能落入其中。

```text
  [tpunetd: 在每个 slice 内启动 ICI fabric — 见 ../ici]
                              │
                              ▼
  ┌─────────────────────────────────────────────────────────────┐
  │ 1. RENDEZVOUS   DiscoverTopologyAndAddressBindings @0x1ccacb80 │
  │    每个 worker 一次 gRPC round trip: GetMultiSliceTopology → │
  │    coordinator 累积 num_slices·num_hosts regs →             │
  │    广播一个 MultiSliceTopologyAndLocation                   │
  └─────────────────────────────────┬───────────────────────────┘
                                    ▼
  ┌─────────────────────────────────────────────────────────────┐
  │ 2. FLEET METADATA   每个进程持有相同的地址表：              │
  │    (slice_id,host_id)→endpoints，每 slice shapes            │
  └─────────────────────────────────┬───────────────────────────┘
                                    ▼
  ┌─────────────────────────────────────────────────────────────┐
  │ 3. BARRIER   xla_tpu_enable_megascale_barrier: 每个 host     │
  │    调用 Barrier(id) → BarrierCoordinator 在 quorum 释放     │
  └─────────────────────────────────┬───────────────────────────┘
                                    ▼
  ┌─────────────────────────────────────────────────────────────┐
  │ 4. STEADY STATE   Send (DCN collectives) + SendHeartBeat    │
  └─────────────────────────────────────────────────────────────┘

  ════ error spine（任一阶段） ════
  ReportError → ErrorReporter @0x1ccb6ea0 → MegascaleErrorAggregator
     → RapidEyeErrorDigestProto + LogErrorDigest（cause classifier）

阶段 1 — Bootstrap rendezvous

runtime 会在 tpunetd 启动 slice 后立即对每个 worker 调用一次 DiscoverTopologyAndAddressBindings。每个 worker 发送恰好一个 GetMultiSliceTopology RPC，携带自己的 (slice_id, host_id, host_addresses, topology_args, incarnation_id)，然后阻塞。coordinator 的 TopologyCoordinator（泛型 Coordinator<Req,Resp,Callback> 模板的一个特化）把注册累积到一个 flat_hash_map<int, SliceState> 中，当 IsComplete() 报告有 num_slices · num_hosts_per_slice 个匹配条目时，它会构建一个 GetMultiSliceTopologyResponse，将其 fan out 给每个 pending callback，并发出通知。这里没有 retry loop；单个每 RPC deadline（--megascale_topology_discovery_timeout）限定整个 rendezvous。内部细节，包括选举、请求 schema、coordinator state machine、字节稳定的响应排序，以及失败路径，都归 Bootstrap section（8 页）所有。

阶段 2 — Fleet metadata

响应会反序列化为实时的 xla::megascale::runtime::MultiSliceTopologyAndLocation C++ class（注意：没有 Proto 后缀；wire form 是 MultiSliceTopologyAndLocationProto）。这是每个下游 consumer 通过 const& 持有的权威 fleet model：Communicator、jellyfish scheduler（ScheduleSendRecvs @ 0x1d6b6520 接收一个 const MultiSliceTopologyAndLocation*），以及 barrier。它回答 fleet 中有什么、我是谁、如何到达一个 peer，以及如何在 fleet 范围内命名芯片。schema，包括 topology model、host identity（incarnation_id）、global addressing、slice shape，以及逐消息 wire decode，归 Fleet Metadata section（9 页）所有。

阶段 3 — 跨主机 barrier

该阶段由 --xla_tpu_enable_megascale_barrier 门控，是执行前 barrier。它复用相同的 Coordinator<> 模板：每个 barrier_id 一个 BarrierCoordinator（由 OnBarrierRequestReceived @ 0x1ccac5c0 惰性插入 backend 的 +0x1b0 map），当 seen-host set 达到 num_workers 时释放所有 waiter。每次调用的 deadline 是 gRPC client deadline（FLAGS_tf_tpu_preexecution_barrier_timeout，默认 30 s）；BarrierRequest 字段 4 是 num_participants，不是 timeout。归 Cross-Host Barrier 所有。

阶段 4 — 稳态

一旦 barrier 清除，地址表会为 DCN collective Send RPC 向 StartTransfer 供给信息，而 StartHeartBeat() @ 0x1ccade60 会开始每 peer 的 SendHeartBeat liveness ping（由 --megascale_use_heartbeat 门控）。只有 coordinator 会记录 "This task is the coordinator. Starting heartbeat with peers."；worker 会对 coordinator heartbeat。heartbeat 失败时，可根据 --megascale_*_restart_* flags 终止或就地重启进程。

错误主干

与所有四个阶段正交：任何检测到故障的主机都会发出一个 ReportError RPC。在 coordinator 上，ErrorReporter::ReportError @ 0x1ccb6ea0 会惰性分配一个 MegascaleErrorAggregator，按 (slice_id:host_id, task_id) 对报告去重，并在 300 ms idle deadline 到达时，或一旦 size() == NumWorkers() 时立即触发 ProcessAndShutdown()。aggregator 运行跨主机 cause classifier（9 值 Cause enum：BAD_TPU_CHIP、NETWORKING_ISSUE、DATA_INPUT_STALL、…），并发出一个 RapidEyeErrorDigestProto。归 Error Aggregator 所有。

章节内容

子系统	页面	拥有内容
Bootstrap rendezvous（8 页）	Bootstrap › Overview	coordinator election、worker registration、topology exchange、convergence、failure handling、tpunetd relationship、ICI handoff
Fleet metadata schema（9 页）	Fleet Metadata › Overview	topology model、host identity、global addressing、ICI-vs-DCN、slice shape、bootstrap exchange、barrier/error usage、field decode
tpunetd wire protocol	tpunetd Protocol	`MegaScaleTransport` / tpunetd daemon wire surface
跨主机 barrier	Cross-Host Barrier	`BarrierCoordinator`、`Coordinator<>` template、quorum 和 timeout
错误聚合	Error Aggregator	`MegascaleErrorAggregator`、RapidEye digest、cause classifier、retention
下方层级	ICI › Overview	Megascale 在 DCN 之上联合起来的 slice 内光学 torus

交叉引用

Bootstrap › Overview — 打开生命周期阶段 1 的跨主机 rendezvous
Fleet Metadata › Overview — rendezvous 广播的 fleet model（阶段 2）
Cross-Host Barrier — 阶段 3，基于同一 Coordinator<> 模板的执行前 barrier
Error Aggregator — 失败主干：ReportError → MegascaleErrorAggregator → RapidEye digest
tpunetd Protocol — daemon 模型的 wire surface
ICI › Overview — Megascale 在 DCN 之上联合的 slice 内光学 torus 层级
ICI › Topology Discovery — 跨越 ICI↔DCN 边界的 TpuTopologyArgsProto 的产生位置

Megascale — 章节地图 ​

摘要 ​

ICI↔DCN 边界 ​

两套坐标系统，在一个点相接 ​

阶段 1 — Bootstrap rendezvous ​

阶段 2 — Fleet metadata ​

阶段 3 — 跨主机 barrier ​

阶段 4 — 稳态 ​

错误主干 ​

章节内容 ​

交叉引用 ​