Megacore 偶/奇拆分
本页中的所有地址、符号和偏移都适用于来自
libtpu-0.0.40-cp314wheel 的libtpu.so(build-id89edbbe81c5b328a958fe628a9f2207d,781,691,048 字节,未剥离符号 - 完整 C++ 符号,构建libtpu_lts_20260413_b_RC00)。.textVMA 等于文件偏移(基址0xe63c000);每个地址都是 VMA。引用的每个符号都存在于完整符号二进制中,并已与 IDA 反编译交叉核对。
摘要
一个 megacore 芯片包含两个物理 TensorCore,运行时可以把它们折叠成 一个 XLA 逻辑设备。twisted-torus all-gather(两阶段 replica-group 构建的 Phase1)必须为每个芯片决定:该芯片贡献 一个 组成员还是 两个;如果是两个,则芯片的每个 core 加入哪个组。本页给出字节精确答案:TwistedTorusND::GetPhase1ReplicaGroups (0x137d3de0) 内部的 Megacore() / CoresPerChip(0) append gate、把 megacore 模式变成组数量乘子的 LogicalDevicesPerChip(0) getter,以及逐 core 的组索引分配 core0 -> 2m(偶数)/ core1 -> 2m+1(奇数)。
反编译确认的决定性结构结论,也与只看循环层面的第一印象相反:偶/奇拆分是 非 megacore 2-core 情况,不是 megacore 情况。megacore 芯片呈现 LogicalDevicesPerChip = 1,因此它把两个物理 core 折叠成单个逻辑参与者,并加入单个组 m(总计 2K 个组)。非 megacore 2-core 芯片呈现 LogicalDevicesPerChip = CoreCount = 2,因此两个 core 扇出到偶/奇组对 {2m, 2m+1}(总计 4K 个组)。Phase0(reduce-scatter)从不拆分 - 它总是把一个芯片的两个 core 共同放入同一个组。
本页是重实现者无法仅从 op 表面或循环结构恢复的三件事的字节精确权威说明:
LogicalDevicesPerChip(0)gate。 对于TpuCoreType=0(TensorCore),Target::LogicalDevicesPerChip(0)(0x1d615b00->TpuTopology::LogicalDevicesPerChip0x20ad3020) 返回Megacore() ? 1 : CoreCount(0)。它是 Phase1 组数量乘子LogicalDevicesPerChip(0) * 2K的唯一来源,也是 megacore 产生2K个组、而非 megacore 2-core 产生4K个组的原因。- Phase1 append gate。 双分支
Megacore() / CoresPerChip(0) != 1测试把每个(m, i, k)三元组路由到组m(单组,两个 core 折叠为一个逻辑设备)或拆分后的{2m, 2m+1}(core0 偶数,core1 奇数)。megacore 分支的拆分条件是逻辑矛盾(CoresPerChip != 1 && CoresPerChip <= 1)- 因此 megacore 芯片总是走单组路径。 - 逐 core 索引。
pair.first(core0)-> 组2m;pair.second(core1)-> 组2m+1。factor-2 LEA(v90 = 2*m,v91 = 2*m+1)在每次外层m迭代中预先计算一次;两个 core 落入相邻组,因此每个 core 的 all-gather 会遍历平面 ICI 链路中互不相交的一半。
Phase0/Phase1 replica-group 驱动器(循环嵌套、组大小、K/2K/R 标量、示例)见 两阶段 Replica-Group 构建;产生 {core0, core1} 对的坐标折叠见 GetReplicaPair3DOnTwistedTorus;产生 K/2K 的形状分类见 形状折叠;工作负载层面的 megacore 折叠(运行时 two-core barrier、MEGACORE(4) 标记、SFLAG slot)见 Megacore Collective Fusion。本页聚焦于组分配中的偶/奇拆分及其背后的 gate。
| Phase1 构建器(拥有该 gate) | TwistedTorusND::GetPhase1ReplicaGroups 0x137d3de0 |
| Phase0 构建器(总是共同分组) | TwistedTorusND::GetPhase0ReplicaGroups 0x137d3560 |
| 组数量乘子 | Target::LogicalDevicesPerChip(0) 0x1d615b00 -> TpuTopology::LogicalDevicesPerChip 0x20ad3020 |
| Megacore 模式 bit | TpuChipConfig::Megacore() 0x20afca00 = byte[TpuChipConfig+0x8] |
| 每芯片 core 数 | Target::CoresPerChip(t) 0x1d615b40 = int32[TpuTopology+0x7c + t*12] |
| Core 数来源 | TpuChipParts::CoreCount(t) 0x20b198e0 (Megacore ? 1 : this) |
| Phase1 组数量 | LogicalDevicesPerChip(0) * 2K - 2K(megacore)或 4K(非 megacore 2-core) |
| 逐 core 索引(拆分) | core0 (pair.first) -> 2m, core1 (pair.second) -> 2m+1 |
| 置信度 | 高 - gate 跳转方向、getter 函数体和 2m/2m+1 LEA 均已反编译验证 |
1. LogicalDevicesPerChip Gate
Phase1 组数量是 LogicalDevicesPerChip(0) * 2K,在 GetPhase1ReplicaGroups 顶部附近计算一次。LogicalDevicesPerChip(0) 是 megacore(2K 个组,无拆分)与非 megacore 2-core(4K 个组,拆分)之间的全部区别。它不会在逐 (m,i,k) append 循环内部读取 - 它只用于确定 vector<ReplicaGroup> 的大小。逐元素拆分决策由 append gate(§2)另行作出,但两者一致:正好在 LogicalDevicesPerChip(0) = 2 时 gate 发生拆分。
Target::LogicalDevicesPerChip - 委托包装器
// Target::LogicalDevicesPerChip(TpuCoreType) — 0x1d615b00
__int64 Target::LogicalDevicesPerChip(Target *this, __int64 core_type) {
return (int)tpu::TpuTopology::LogicalDevicesPerChip(
*(_QWORD *)(this + 952), // Target+0x3b8 == the TpuTopology*
core_type);
}
```text
`Target+0x3b8`(十进制 `952`)是整个 twist 子系统读取的 `TpuTopology*`;同一字段也供给 `CoresPerChip` 和 `Megacore()` gate(§2)。Phase1 调用它时使用 `core_type = 0`(TensorCore)。
### `TpuTopology::LogicalDevicesPerChip` - megacore 在这里变成 1
```c
// tpu::TpuTopology::LogicalDevicesPerChip(TpuCoreType) — 0x20ad3020 (condensed)
__int64 TpuTopology::LogicalDevicesPerChip(__int64 self, __int64 core) {
cc = TpuChipParts::CoreCount(*(_QWORD *)(self + 8), core);
if (!cc) return 0; // core type absent on this chip
if (core != 2) { // TensorCore (0) path
if (core != 0) FATAL("Unsupported core type"); // tpu_topology.cc:538
if (TpuChipConfig::Megacore(*(_QWORD *)(self + 24)))
return 1; // MEGACORE: 1 logical device
return TpuChipParts::CoreCount(*(_QWORD *)(self + 8), 0); // else: CoreCount
}
/* core == 2 (SparseCore) path: Megachip / SharedMemoryCount-based — out of scope */
}TensorCore 分支(core == 0)精确化简为:
LogicalDevicesPerChip(0) = Megacore() ? 1 : CoreCount(0)
```text
`TpuChipConfig::Megacore()` (`0x20afca00`) 是一次单字节读取:
```c
// TpuChipConfig::Megacore() — 0x20afca00
__int64 TpuChipConfig::Megacore(TpuChipConfig *this) {
return *((unsigned __int8 *)this + 8); // byte[TpuChipConfig+0x8]
}因此,带两个物理 TensorCore 的 megacore 芯片返回 1(两个 core 是一个逻辑设备);非 megacore 芯片返回 CoreCount(0)(2-core 芯片为 2,single-core 芯片为 1)。
陷阱 -
TpuTopology::LogicalDevicesPerChip的 SparseCore 分支(core == 2)使用不同公式(受Megachip()和SharedMemoryCount控制),不是Megacore ? 1 : CoreCount。Phase1 只会调用LogicalDevicesPerChip(0)(TensorCore),所以本页记录的是 TensorCore 拆分。SparseCore collective 的 Phase1 组是否按同一规则拆分,并未由TwistedTorusND执行 - 它位于sparse_core::collectivetopology builder 中。见 SC 侧 Twist。
组数量乘法
// GetPhase1ReplicaGroups — 0x137d3de0, prologue (decompiled, condensed)
v18 = *((_QWORD *)a2 + 190); // 2K (max_dim, [obj+0x5f0])
v93 = Target::CoresPerChip(a3, 0); // cores_per_chip (held for the gate)
v20 = v18 * Target::LogicalDevicesPerChip(a3, 0); // group count = 2K · LDPC(0)
v23 = operator new(48 * v20); // vector<ReplicaGroup>(group_count)
```text
`v18` 是 `2K`;`Target::LogicalDevicesPerChip(a3, 0)` 是 `1`(megacore)或 `2`(非 megacore 2-core);乘积 `v20` 是组数量 `2K` 或 `4K`。`48 * v20` 分配(`ReplicaGroup` 为 48 字节)是唯一读取 `LogicalDevicesPerChip` 的位置 - 这确认了*数量*由乘子设置,而*路由*由逐元素 gate(§2)设置。如果重实现者只按平面范围分配(没有 `LDPC` 因子),在非 megacore 2-core 情况下会少分配 `2x`,并让组 `2m+1` 索引越界。
| 模式 | `Megacore()` | `CoreCount(0)` | `LogicalDevicesPerChip(0)` | Phase1 组数量 |
|---|---|---|---|---|
| megacore | `1` | `2` | `1` | `2K` |
| 非 megacore,1 core/chip | `0` | `1` | `1` | `2K` |
| 非 megacore,2 cores/chip | `0` | `2` | `2` | `4K` |
---
## 2. Phase1 Append Gate - 字节精确
逐元素路由是在 `Megacore()` 和 `CoresPerChip(0)` 上的双分支测试,在坐标折叠返回 `{pair.first = core0, pair.second = core1}` 后,对每个 `(m, i, k)` 三元组求值一次。该 gate 是 `0x137d4348..0x137d4462` 的实时字节序列;反编译形式是最清楚的陈述。
### 反编译
```c
// GetPhase1ReplicaGroups — 0x137d3de0, the append gate (decompiled, condensed)
// per-m precompute: v90 = 2*m (even), v91 = 2*m+1 (odd), v94 = 48*m (single group m)
// per-(m,i,k): ReplicaPair3DOnTwistedTorus = pair.first (core0); v43 = pair.second (core1)
if ( TpuChipConfig::Megacore(*(_QWORD *)(*((_QWORD *)v88 + 119) + 24)) ) { // Target+0x3b8 -> TpuChipConfig
if ( v93 != 1 && *(int *)(*((_QWORD *)v88 + 119) + 124) <= 1 ) // CoresPerChip(0)!=1 && [+0x7c]<=1
goto SPLIT; // (DEAD: contradiction, see below)
// megacore, no split -> fall through to SINGLE group m
} else { // non-megacore
if ( v93 != 1 ) // CoresPerChip(0) != 1
goto SPLIT; // 2-core chip -> split
// non-megacore single-core -> fall through to SINGLE group m
}
// ---- SINGLE: append core0 (pair.first) into group m (offset 48*m) ----
groups[m].add(pair.first); // -> group m
goto DONE;
SPLIT:
groups[2*m ].add(pair.first ); // core0 -> group 2m (even)
groups[2*m+1].add(pair.second); // core1 -> group 2m+1 (odd)
DONE:v88 是 Target*;*((_QWORD *)v88 + 119) 是 Target+0x3b8(119*8 = 952 = 0x3b8),即 TpuTopology*;+24 到达 TpuChipConfig 以读取 Megacore();+124 是 TpuTopology+0x7c,即内联读取的 CoresPerChip(0)。v93 是已在 prologue 中加载的同一个 CoresPerChip(0) 值。
真值表
Megacore() | CoresPerChip(0) | 采取的分支 | 路由 |
|---|---|---|---|
1(megacore) | 2 | megacore split test 2 != 1 && 2 <= 1 = false -> fall through | 单个组 m |
1(megacore) | 1 | megacore split test 1 != 1 && … = false -> fall through | 单个组 m |
0(非 megacore) | 1 | 1 != 1 = false -> fall through | 单个组 m |
0(非 megacore) | 2 | 2 != 1 = true -> SPLIT | 拆分 core0->2m, core1->2m+1 |
怪异点 - megacore split 分支是死分支。 megacore 分支的拆分条件是
CoresPerChip(0) != 1 && CoresPerChip(0) <= 1,对任何整数都不可能为真(v93正是同一个CoresPerChip(0))。因此无论有多少物理 core,megacore 芯片都总是到达 single-group-m路径。该条件在结构上存在(编译器发出了cmp [rbp-0xb0],1; je 0x4462和cmp [TpuTopology+0x7c],1; jg 0x4462两个跳转),但只有非 megacore 分支的CoresPerChip(0) != 1测试能真正路由到 split。这是在字节层面确认 megacore 从不拆分:不是 megacore 选择 single-group,而是唯一可达的 split 路径属于非 megacore。
为什么拆分把 core0 路由到偶数组、core1 路由到奇数组
两个 factor-2 组索引在每个外层 m 中预先计算一次(v90 = 2*m,v91 = 2*m+1),因此 split append 是一对固定偏移 48*2m 和 48*(2m+1)。坐标折叠返回芯片的 {core0, core1} device-id 对;拆分把 core0 送到偶数组,把 core1 送到奇数组。因此,非 megacore 2-core 芯片的两个 core 会在互不相交的组半边上执行 all-gather:每个偶数组 2m 只携带切片 m 中每个芯片的 core0,每个奇数组 2m+1 只携带 core1。每个 core 的 all-gather 流量使用平面 ICI 链路的不同切片,而不是两个 core 竞争同一个组,这会在两个逻辑设备之间平衡带宽。在 megacore 模式中,芯片是一个逻辑设备,因此每芯片只有一个 all-gather 参与者和一个组 m - 跨两个物理 core 的工作负载拆分发生在该单一逻辑设备内部,由运行时折叠完成(Megacore Collective Fusion),而不是由 replica-group 构建完成。
3. Phase0 总是共同分组 - 对比
Phase0(沿 2K ring 的 reduce-scatter)运行相同的 Megacore() / CoresPerChip(0) 测试,但它的两个 append - 当两者都触发时 - 写入同一个组,从不写入偶/奇组对。Phase1 执行的拆分在 Phase0 中没有对应物。
// GetPhase0ReplicaGroups — 0x137d3560, the second-core gate (decompiled, condensed)
// first append already done: groups[g].add(pair.first) // core0, group offset 48*g
if ( TpuChipConfig::Megacore(*(_QWORD *)(*((_QWORD *)v86 + 119) + 24)) ) {
if ( v88 == 1 || *(int *)(*((_QWORD *)v86 + 119) + 124) > 1 ) // CoresPerChip(0)==1 || [+0x7c]>1
goto SKIP_SECOND; // megacore 2-core: skip core1
} else {
if ( v88 == 1 ) // single-core: skip core1
goto SKIP_SECOND;
}
groups[g].add(pair.second); // core1 -> SAME group g (offset 48*g, identical to core0's)
SKIP_SECOND:
```text
`v88` 是 `CoresPerChip(0)`。第二个 append(`pair.second`,core1)目标是与第一个(`pair.first`,core0)相同的组偏移 `48*g` - 两者都写入由单个索引 `g = k*R + i` 计算出的组,该索引不依赖 core。因此 reduce-scatter ring 让芯片的两个 core **共驻**在同一个 ring 中;它们直到 all-gather 才扇出。
Phase0 gate 在两种情况下跳过第二个 core append:
| `Megacore()` | `CoresPerChip(0)` | 第二次 append? | 原因 |
|---|---|---|---|
| `1`(megacore) | `2` | **跳过**(`CoresPerChip(0) > 1`) | 一个逻辑设备 - `pair.first` 已覆盖它 |
| `1`(megacore) | `1` | 跳过(`== 1`) | 单 core |
| `0`(非 megacore) | `2` | **append** | 两个 core 都加入 ring group |
| `0`(非 megacore) | `1` | 跳过(`== 1`) | 单 core |
> **说明 -** 这种不对称正是关键。在 megacore 模式中,reduce-scatter 组只 append `pair.first`(单个逻辑设备的 core0 id),因为两个物理 core 在 ring 上是*一个*设备;all-gather 同样使用单个组 `m`。在非 megacore 2-core 模式中,reduce-scatter 把*两个* core id 共同放入一个 ring(它们共享芯片的 ICI ring 位置,因此分离它们没有带宽收益),但 all-gather 把它们拆分到偶/奇半边(在那里分离它们*确实*能平衡正交平面的链路)。因此 Phase0 和 Phase1 读取相同输入,却得到相反的组形状:RS 共同分组,AG 拆分。
---
## 4. 三个拓扑 Getter
该 gate 通过 `Target` 对象位于 `Target+0x3b8` 的 `TpuTopology*` 读取三个运行时拓扑值。三者都是小型叶函数;重实现者必须精确复现它们,因为 gate 的正确性依赖于 §2 中的 contradiction-and-fallthrough 结构,而该结构又依赖这些 getter 返回正确字段。
### `Target::CoresPerChip(t)` - `0x1d615b40`
```c
// Target::CoresPerChip(TpuCoreType) — 0x1d615b40
__int64 Target::CoresPerChip(Target *this, unsigned int core_type) {
if (core_type >= 3) BUG();
return *(int *)(*(_QWORD *)(this + 952) + 12LL * core_type + 124);
// Target+0x3b8 (TpuTopology*) + 0x7c + core_type·12
}CoresPerChip(0) 是 int32[TpuTopology + 0x7c] - 一个按 TpuCoreType 索引、步长 12 字节的数组,表示芯片的物理 core 数(不是逻辑数)。对于 megacore 2-core 芯片,即使 LogicalDevicesPerChip(0) = 1,这里仍为 2;这种不匹配(CoresPerChip = 2,LDPC = 1)正是 gate 区分的对象。
TpuChipConfig::Megacore() - 0x20afca00
byte[TpuChipConfig + 0x8](§1)。TpuChipConfig 通过 Target+0x3b8(TpuTopology*)+ 0x18(TpuChipConfig*)到达。
TpuTopology::LogicalDevicesPerChip(0) - 0x20ad3020
Megacore() ? 1 : CoreCount(0)(§1),由 Target::LogicalDevicesPerChip(0x1d615b00)包装。CoreCount 通过 TpuChipParts::CoreCount(t)(0x20b198e0)解析。
| Getter | 地址 | 返回 | 字段 |
|---|---|---|---|
Target::CoresPerChip(t) | 0x1d615b40 | 类型 t 的物理 core 数 | int32[TpuTopology+0x7c + t*12] |
TpuChipConfig::Megacore() | 0x20afca00 | megacore 模式 bit | byte[TpuChipConfig+0x8] |
Target::LogicalDevicesPerChip(t) | 0x1d615b00 | (包装器) | 委托给 0x20ad3020 |
TpuTopology::LogicalDevicesPerChip(t) | 0x20ad3020 | 每芯片逻辑设备数 | Megacore ? 1 : CoreCount(t)(TC 分支) |
TpuChipParts::CoreCount(t) | 0x20b198e0 | 类型 t 的 core 数 | per-TpuChipParts 字段 |
5. Megacore 拆分表 - 组数量与路由
对 TensorCore Phase1 构建,结合 §1 的组数量乘子与 §2 的 append gate 后,完整决策如下:
| 模式 | Megacore() | CoresPerChip(0) | LogicalDevicesPerChip(0) | Phase1 组 | 逐 (m,i,k) 路由 |
|---|---|---|---|---|---|
| megacore | 1 | 2 | 1 | 2K | 组 m(两个 core -> 一个逻辑设备) |
| 非 megacore,1 core | 0 | 1 | 1 | 2K | 组 m(单 core) |
| 非 megacore,2 cores | 0 | 2 | 2 | 4K | core0 -> 组 2m(偶数),core1 -> 组 2m+1(奇数) |
与 Phase0 对比一起读:Phase0 组数量始终为 K*R(它使用 CoresPerChip,不是 LogicalDevicesPerChip,并且从不翻倍);当两个 core 都被 append 时,它们落入同一个 Phase0 组。2K 与 4K 的翻倍是 Phase1 专属的 LogicalDevicesPerChip(0) 结果,而偶/奇路由是 Phase1 专属的非 megacore CoresPerChip(0) != 1 分支结果。匹配的组乘积和 K, K, 2K 示例见 两阶段 Replica-Group 构建 §4-§5。
陷阱 -
Megacore()测试确实出现在 split guard 的第一个分支中,但 megacore 分支的拆分条件是矛盾,永不触发:megacore 始终使用单组m(2K个组)。偶/奇拆分只会通过非 megacore 的CoresPerChip(0) != 1分支到达(LogicalDevicesPerChip(0) = CoreCount = 2 -> 4K个组)。不要把 guard 读成“megacore: core0 -> even / core1 -> odd”。
6. 函数映射
| 函数 | 地址 | 作用 |
|---|---|---|
TwistedTorusND::GetPhase1ReplicaGroups | 0x137d3de0 | 拥有偶/奇 append gate;组数量 LDPC(0)*2K |
TwistedTorusND::GetPhase0ReplicaGroups | 0x137d3560 | 总是把两个 core 共同分到同一个组 |
Target::LogicalDevicesPerChip | 0x1d615b00 | 包装器 -> TpuTopology::LogicalDevicesPerChip |
TpuTopology::LogicalDevicesPerChip | 0x20ad3020 | Megacore ? 1 : CoreCount(TC 分支)- gate 的来源 |
TpuChipConfig::Megacore | 0x20afca00 | byte[TpuChipConfig+0x8] megacore 模式 bit |
Target::CoresPerChip | 0x1d615b40 | int32[TpuTopology+0x7c + t*12] 物理 core 数 |
TpuChipParts::CoreCount | 0x20b198e0 | 按 core 类型的物理 core 数 |
GetReplicaPair3DOnTwistedTorus | 0x1c893400 | 产生该 gate 路由的 {core0, core1} 对 |
7. 未解决事项
- 每芯片
CoreCount(0)字面值的生成。LogicalDevicesPerChip(0) = Megacore ? 1 : CoreCount(0)是字节精确的;字面CoreCount(0)值(megacore/2-core TPU 为 2,single-core 为 1)来自由 chip-config proto 填充的TpuChipParts。getter 公式已确认;按 codename 的常量是 proto 依赖,此处未提取。中等。 - SparseCore Phase1 拆分。 Phase1 只读取
LogicalDevicesPerChip(0)/CoresPerChip(0)(TensorCore,t=0)。SparseCore collective 的 Phase1 组是否按同一个Megacore ? 1 : CoreCount规则拆分(0x20ad3020的core == 2分支是不同的、基于Megachip/SharedMemoryCount的公式)并未由TwistedTorusND执行。见 SC 侧 Twist。低。 arg >= 1多 shard 路径。 该 gate 是针对实时arg == 0single-phase collective 解码的;坐标折叠的arg >= 1入口在 shard gate 后被 CHECK 标记为不可达(GetPerColorShardIdTable,两阶段 Replica-Group 构建 §6)。未来 multi-shard 构建是否改变逐 core 路由尚未执行验证。低。
交叉引用
Twist 算法(本节)
- Twisted Torus - 章节地图 - 子系统地图;引用本页说明由
LogicalDevicesPerChip控制的拆分 - 两阶段 Replica-Group 构建 - Phase0/Phase1 驱动器、组大小,以及
K/2K/R标量(链接到这里获取字节精确 gate) - GetReplicaPair3DOnTwistedTorus - 产生该 gate 路由的
{core0, core1}对的坐标折叠 - 形状折叠 -
K、2K和R = num-2K-axes >= 2 ? 2K : K平面维度的来源 - TwistedTorusND::BuildStrategy - 与这些 device-id 列表互补的逐 color ring-neighbour 发射侧
- SC 侧 Twist - SparseCore twisted-torus topology builder(
core == 2LogicalDevicesPerChip分支)
同级章节
- Megacore Collective Fusion - 工作负载层面的 megacore 折叠(运行时 two-core barrier、
MEGACORE(4)标记、SFLAG slot)- megacore 芯片使用的逻辑设备内部拆分,而不是偶/奇组拆分 - 物理 Core 放置 - 逻辑设备如何映射到物理 core
- 返回索引