Internal Wiki

主题

逐字段解码 ​

Megascale 机群元数据的完整 wire schema，恢复自 libtpu.so 内部的 protodesc_cold FileDescriptorProto 池。除非另有说明，所有消息都属于 package xla.megascale.runtime；topology-args 系列属于 package tpu。

SliceInfo ​

platforms/xla/megascale/common/runtime.proto — 每个 slice 的机群条目。

protobuf
message SliceInfo {
  int32                    slice_id          = 1;   // fleet-wide slice index
  // field 2 retired (gap)
  tpu.TpuTopologyArgsProto tpu_topology_args = 3;   // slice 3D shape
}
```text

这是规范的 `SliceInfo`：slice 索引加形状，主机端点则单独保存在 `NetworkAddressMapping` 列表中。

## MultiSliceTopologyAndLocationProto

`runtime.proto` — 序列化的 FLEET 对象（`MultiSliceTopologyAndLocation` 类的 wire 形式）。

```protobuf
message MultiSliceTopologyAndLocationProto {
  int32              local_slice_id = 1;   // THIS process's slice
  int32              local_host_id  = 2;   // THIS process's host (in slice)
  repeated SliceInfo slice_info     = 3;   // every slice in the fleet
  int64              incarnation_id = 4;   // fleet generation token
}

由 Communicator::Create、MegaScaleMultiSliceConfig::Create 和 AotMegaScaleMultiSliceConfig::Create 接受。local_* 字段使其能够自定位。

MultiSliceTopologyInfo ​

platforms/xla/megascale/runtime/communication/transport.proto — 协调器序列化到 bootstrap 响应中的 wire 载荷。

protobuf
message MultiSliceTopologyInfo {
  repeated SliceInfo             slice_info       = 1;   // all slices
  repeated NetworkAddressMapping address_mappings = 2;   // all host endpoints
  int64                          incarnation_id   = 3;
}
```text

将 slice 清单与完整端点表配对。序列化到 `GetMultiSliceTopologyResponse.serialized_topology_info` 中。

## NetworkAddressMapping / HostNetworkAddress / EndpointAddresses

`third_party/tensorflow/compiler/xla/megascale/addresses.proto` — DCN 可达性层。

```protobuf
message HostNetworkAddress {
  string address                 = 1;   // ip:port / host:port of a NIC
  string interface_name          = 2;   // NIC name (DCN port)
  int32  numa_node               = 4;   // NUMA affinity of this NIC
  string host_name_for_debugging = 3;   // human-readable host label
}
message NetworkAddressMapping {
  int32  slice_id                = 1;   // owning slice
  int32  host_id                 = 2;   // host within slice
  repeated HostNetworkAddress addresses = 3;   // one per DCN interface
}
message EndpointAddresses {
  repeated NetworkAddressMapping address_mappings = 1;   // whole fleet
}

(slice_id, host_id) 是通用主机键。注意 HostNetworkAddress 中不连续的字段编号 (1, 2, 4, 3)；一个已停用字段留下了空洞。

GetMultiSliceTopologyRequest / Response ​

transport.proto — bootstrap RPC 载荷。

protobuf
message GetMultiSliceTopologyRequest {
  NetworkAddressMapping    address_mapping   = 1;   // THIS host's identity + endpoints
  tpu.TpuTopologyArgsProto tpu_topology_args = 2;   // THIS slice's shape
  int64                    incarnation_id    = 3;   // THIS process's generation
}
message GetMultiSliceTopologyResponse {
  bytes serialized_topology_info = 1;   // serialized MultiSliceTopologyInfo
}
```text

请求的 `slice_id` / `host_id` 位于 `address_mapping` *内部*；响应是单个不透明的 `bytes` blob。

## TpuTopologyArgsProto

`learning/45eac/tpu/runtime/topology/tpu_topology_args.proto` (package
`tpu`, editions) — slice 形状描述符。完整解码见
[Slice 形状](slice-shape.md) 页面。

```protobuf
message TpuTopologyArgsProto {
  tpu.TpuVersionProto      version                  = 1;
  string                   variant                  = 2;
  tpu.TpuPlatformTypeProto platform_type            = 3;
  string                   chip_config_name         = 4;
  tpu.TpuDimensionsProto   chips_per_host_bounds    = 5;
  tpu.TpuDimensionsProto   host_bounds              = 6;
  tpu.TpuWrapProto         wrap                     = 7;
  bool                     twist                    = 8;
  bool                     enhanced_barrier_enabled = 9;
  bool                     use_continuations        = 11;
  tpu.TpuRoutingStrategyProto routing_strategy       = 12;
  SubSlice                 sub_slice                = 10;
  message SubSlice {
    tpu.TpuDimensionsProto chips_per_host_bounds    = 1;
    tpu.TpuDimensionsProto host_bounds              = 2;
  }
}

DCNTopology ​

third_party/tensorflow/compiler/xla/megascale/dcn_topology.proto — 跨 slice reduction 计划。

protobuf
message DCNTopology {
  oneof representation {                       // "representation"
    SymmetricTree symmetric_tree = 1;
    TreeNode      tree           = 2;
  }
  message SymmetricTree {
    repeated int32 branching_per_layer = 1;    // balanced tree fan-out per layer
  }
  message TreeNode {
    string            label             = 1;   // optional node label
    repeated TreeNode nodes             = 2;   // children
    repeated SliceRange slice_ranges    = 3;   // leaf slice ranges
    float             egress_constraint = 4;   // per-node egress cap
    bool              ring_transfers    = 5;   // ring vs tree at this node
    message SliceRange { int32 slice_id_start = 1; int32 slice_id_end = 2; }
  }
}
```text

`SymmetricTree` 是紧凑形式；`TreeNode` 是显式形式，其叶子命名连续的 `[slice_id_start, slice_id_end]` 范围。

## TpuCoreLocationProto

`platforms/xla/megascale/runtime/common/core_location.proto` — 每个 core 的句柄（由 `MultiSliceTpuCoreLocation` 类与其 slice 关联）。

```protobuf
message TpuCoreLocationProto {
  tpu.TpuCoreTypeProto core_type         = 1;   // TENSOR / BARNA / SPARSE
  int32                logical_device_id = 2;   // device id within its slice
}
enum tpu.TpuCoreTypeProto {
  TPU_CORE_TYPE_INVALID=0; TPU_CORE_TYPE_TENSOR_CORE=1;
  TPU_CORE_TYPE_BARNA_CORE=2; TPU_CORE_TYPE_SPARSE_CORE=3;
}

PerSliceTpuDimensionsProto / Topology ​

runtime.proto — 轻量级的每 slice (x,y,z) 视图，用于不需要完整形状的地方（例如 MEGASCALE_TOPOLOGY 覆盖路径）。

protobuf
message PerSliceTpuDimensionsProto {
  int32 slice_id = 1; int32 x = 2; int32 y = 3; int32 z = 4;
}
message Topology { repeated PerSliceTpuDimensionsProto slices = 1; }
```text

## 辅助维度和枚举类型

```protobuf
message TpuDimensionsProto   { int32 x=1; int32 y=2; int32 z=3; int32 w=4; }
message TpuWrapProto         { bool  x=1; bool  y=2; bool  z=3; }
message TpuDegradedAxesProto { bool  x=1; bool  y=2; bool  z=3; }

enum TpuRoutingStrategyProto { ROUTING_DEFAULT=0; ROUTING_MESH=1; ROUTING_NHOP=2; }
enum TpuVersionProto {
  TPU_VERSION_INVALID=0; TPU_VERSION_JELLYFISH=1; TPU_VERSION_DRAGONFISH=2;
  TPU_VERSION_PUFFERFISH=3; TPU_VERSION_VIPERFISH=4; TPU_VERSION_GHOSTLITE=5;
  TPU_VERSION_6acc60406=6;
}
enum TpuPlatformTypeProto {
  TPU_PLATFORM_TYPE_INVALID=0; TPU_PLATFORM_TYPE_HARDWARE=1;
  TPU_PLATFORM_TYPE_GRM=2; TPU_PLATFORM_TYPE_ISS=3;
}
message TpuConfiguredPropertiesProto {
  tpu.TpuDegradedAxesProto    degraded_axes           = 1;
  bool                        is_nhop_source_relative = 2;
  tpu.TpuRoutingStrategyProto routing_strategy        = 3;
}

运行时选项和 collective slack（辅助） ​

runtime.proto 还携带两个机群作用域的消息：

protobuf
message MegascaleRuntimeOptions {                 // all fields optional (oneof presence)
  bool abort_on_errors      = 1;
  bool abort_on_hangs       = 2;
  bool debug_dump_on_error  = 3;
}
message CollectiveSlackInformation {
  repeated BudgetInfo     default_budgets         = 1;
  repeated CollectiveInfo collectives             = 2;
  CollectiveInfo          default_collective_info = 3;
  message BudgetInfo {
    repeated google.protobuf.Duration transmission_budget_per_level = 1;
    google.protobuf.Duration          safety_margin                 = 2;
    bool                              add_random_delay              = 3;
    repeated int32                    levels_with_delay             = 4;
    repeated google.protobuf.Duration explicit_delays               = 5;
    int64                             min_size_for_delay            = 6;
  }
  message CollectiveInfo {
    string                   key                  = 1;
    google.protobuf.Duration slack                = 2;
    oneof budget {
      BudgetInfo             transmission_budget   = 3;
      int32                  default_budget_index  = 4;
    }
  }
}
```text

`MegascaleRuntimeOptions` 映射 `--megascale_*_abort_on_*` 标志；`CollectiveSlackInformation` 调整每个 collective 的 DCN 传输预算（每个 tree level 的时长），并与机群视图一起传入 `HostCommandSchedulerFactory`。

## 在 libtpu.so 中验证存在

| 类型 / 字符串 | 出现次数 (`rg -a -o <s> \| wc -l`) |
|---------------|-------------|
| `MultiSliceTopologyAndLocationProto` (字符串) | 37 |
| `MultiSliceTopologyAndLocation` (不同的成员函数) | 7 |
| `NetworkAddressMapping` | 102 |
| `HostNetworkAddress` | 58 |
| `EndpointAddresses` | 136 |
| `TpuTopologyArgsProto` | 82 |
| `PerSliceTpuDimensionsProto` | 19 |
| `chips_per_host_bounds` | 25 |
| `host_bounds` | 21 |
| `chip_config_name` | 13 |
| `serialized_topology_info` | 1 |
| `SetGlobalChipId` | 161 |
| `GetChipCoordinates` | 107 |
| `ChipLocationToCoordinate` | 32 |
| `global_core_id` / `global_chip_id` | 8 / 6 |
| `GetSliceLocalDeviceId` / `ToSliceAndLogicalDeviceId` | 2 / 1 |

## 交叉引用

- [机群元数据 › 概览](overview.md) — 这些 wire 字段填充的模型
- [拓扑模型](topology-model.md) — 解码后的坐标如何组装为 `ToroidalTopology`
- [全局寻址](global-addressing.md) — 此处解码的 `global_core_id` / `global_chip_id` 推导
- [Slice 形状](slice-shape.md) — 每 slice 边界字段及其版本控制门控
- [Bootstrap 交换](bootstrap-exchange.md) — 通过 DCN 携带此 blob 的 rendezvous
- [ICI › 拓扑发现](../../ici/topology-discovery.md) — 每 slice `tpu_topology_args.proto` 的 slice 内生产者