Skip to content

Megacore 偶/奇拆分

本页中的所有地址、符号和偏移都适用于来自 libtpu-0.0.40-cp314 wheel 的 libtpu.so(build-id 89edbbe81c5b328a958fe628a9f2207d,781,691,048 字节,未剥离符号 - 完整 C++ 符号,构建 libtpu_lts_20260413_b_RC00)。.text VMA 等于文件偏移(基址 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::LogicalDevicesPerChip 0x20ad3020) 返回 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)-> 组 2mpair.second(core1)-> 组 2m+1。factor-2 LEA(v90 = 2*mv91 = 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 模式 bitTpuChipConfig::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 - 委托包装器

c
// 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)精确化简为:

c
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::collective topology builder 中。见 SC 侧 Twist

组数量乘法

c
// 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:

v88Target**((_QWORD *)v88 + 119)Target+0x3b8119*8 = 952 = 0x3b8),即 TpuTopology*+24 到达 TpuChipConfig 以读取 Megacore()+124TpuTopology+0x7c,即内联读取的 CoresPerChip(0)v93 是已在 prologue 中加载的同一个 CoresPerChip(0) 值。

真值表

Megacore()CoresPerChip(0)采取的分支路由
1(megacore)2megacore split test 2 != 1 && 2 <= 1 = false -> fall through单个m
1(megacore)1megacore split test 1 != 1 && … = false -> fall through单个m
0(非 megacore)11 != 1 = false -> fall through单个m
0(非 megacore)22 != 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 0x4462cmp [TpuTopology+0x7c],1; jg 0x4462 两个跳转),但只有非 megacore 分支的 CoresPerChip(0) != 1 测试能真正路由到 split。这是在字节层面确认 megacore 从不拆分:不是 megacore 选择 single-group,而是唯一可达的 split 路径属于非 megacore。

为什么拆分把 core0 路由到偶数组、core1 路由到奇数组

两个 factor-2 组索引在每个外层 m 中预先计算一次(v90 = 2*mv91 = 2*m+1),因此 split append 是一对固定偏移 48*2m48*(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 中没有对应物。

c
// 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 = 2LDPC = 1)正是 gate 区分的对象。

TpuChipConfig::Megacore() - 0x20afca00

byte[TpuChipConfig + 0x8](§1)。TpuChipConfig 通过 Target+0x3b8TpuTopology*+ 0x18TpuChipConfig*)到达。

TpuTopology::LogicalDevicesPerChip(0) - 0x20ad3020

Megacore() ? 1 : CoreCount(0)(§1),由 Target::LogicalDevicesPerChip0x1d615b00)包装。CoreCount 通过 TpuChipParts::CoreCount(t)0x20b198e0)解析。

Getter地址返回字段
Target::CoresPerChip(t)0x1d615b40类型 t 的物理 core 数int32[TpuTopology+0x7c + t*12]
TpuChipConfig::Megacore()0x20afca00megacore 模式 bitbyte[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) 路由
megacore1212Km(两个 core -> 一个逻辑设备)
非 megacore,1 core0112Km(单 core)
非 megacore,2 cores0224Kcore0 -> 组 2m(偶数),core1 -> 组 2m+1(奇数)

与 Phase0 对比一起读:Phase0 组数量始终为 K*R(它使用 CoresPerChip,不是 LogicalDevicesPerChip,并且从不翻倍);当两个 core 都被 append 时,它们落入同一个 Phase0 组。2K4K 的翻倍是 Phase1 专属的 LogicalDevicesPerChip(0) 结果,而偶/奇路由是 Phase1 专属的非 megacore CoresPerChip(0) != 1 分支结果。匹配的组乘积和 K, K, 2K 示例见 两阶段 Replica-Group 构建 §4-§5。

陷阱 - Megacore() 测试确实出现在 split guard 的第一个分支中,但 megacore 分支的拆分条件是矛盾,永不触发:megacore 始终使用单组 m2K 个组)。偶/奇拆分只会通过非 megacoreCoresPerChip(0) != 1 分支到达(LogicalDevicesPerChip(0) = CoreCount = 2 -> 4K 个组)。不要把 guard 读成“megacore: core0 -> even / core1 -> odd”。


6. 函数映射

函数地址作用
TwistedTorusND::GetPhase1ReplicaGroups0x137d3de0拥有偶/奇 append gate;组数量 LDPC(0)*2K
TwistedTorusND::GetPhase0ReplicaGroups0x137d3560总是把两个 core 共同分到同一个组
Target::LogicalDevicesPerChip0x1d615b00包装器 -> TpuTopology::LogicalDevicesPerChip
TpuTopology::LogicalDevicesPerChip0x20ad3020Megacore ? 1 : CoreCount(TC 分支)- gate 的来源
TpuChipConfig::Megacore0x20afca00byte[TpuChipConfig+0x8] megacore 模式 bit
Target::CoresPerChip0x1d615b40int32[TpuTopology+0x7c + t*12] 物理 core 数
TpuChipParts::CoreCount0x20b198e0按 core 类型的物理 core 数
GetReplicaPair3DOnTwistedTorus0x1c893400产生该 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 规则拆分(0x20ad3020core == 2 分支是不同的、基于 Megachip/SharedMemoryCount 的公式)并未由 TwistedTorusND 执行。见 SC 侧 Twist。低。
  • arg >= 1 多 shard 路径。 该 gate 是针对实时 arg == 0 single-phase collective 解码的;坐标折叠的 arg >= 1 入口在 shard gate 后被 CHECK 标记为不可达(GetPerColorShardIdTable两阶段 Replica-Group 构建 §6)。未来 multi-shard 构建是否改变逐 core 路由尚未执行验证。低。

交叉引用

Twist 算法(本节)

同级章节

  • Megacore Collective Fusion - 工作负载层面的 megacore 折叠(运行时 two-core barrier、MEGACORE(4) 标记、SFLAG slot)- megacore 芯片使用的逻辑设备内部拆分,而不是偶/奇组拆分
  • 物理 Core 放置 - 逻辑设备如何映射到物理 core
  • 返回索引