Skip to content

TpuVersion 感知的 flag 前缀分发

本页中的所有地址、符号和计数都适用于 libtpu-0.0.40-cp314 wheel 中的 libtpu.so(构建 libtpu_lts_20260413_b_RC00,build-id md5 89edbbe81c5b328a958fe628a9f2207d,ELF x86-64 DYN,未 strip)。其他版本会有所不同。

摘要

libtpu 携带七个可识别的按子系统划分的 flag 前缀命名空间:四个 TPU 世代 codename xla_jf_*(Jellyfish)、xla_pf_*(Pufferfish)、xla_vf_*(Viperfish)、xla_gf_*(6acc60406/v7x),再加上 xla_sc_*(SparseCore)、barna_core_*(embedding engine)以及通用的 xla_tpu_* 主表面。重实现者自然会问:“当活动世代是 Viperfish 时,xla_jf_* flag 会不会 no-op?”本页从二进制中恢复答案,而答案是不会,并且这个问题本身形式不对:flag 注册和 flag 应用都没有世代门控。flag 名称中的 codename 是静态的作者命名空间,不是运行时路由键。

参照框架是 flag-families 路由图,它建立了 prefix → sink 分类(通用 xla_*DebugOptionsxla_tpu_* / codename / xla_sc_* / barna_core_*TpuCompilationEnvironment(TCE);megascale_* / tpu_* → 独立 runtime)。本页是其下层:实际的分发。每个路由到 TCE 的 flag,包括全部四个 codename 家族,都会无条件注册为一个 absl::Flag<T>,通过单个全局 flag↔field 哈希表(FlagFieldMappings)1:1 绑定到一个 TCE proto 字段,并由一个反射循环应用;该循环遍历 TCE 描述符的每个字段,且其签名中任何位置都没有 TpuVersion 参数。在该循环中,xla_jf_* flag 和 xla_vf_* flag 是平等成员;它们唯一的差别是各自写入哪个 TCE 字段以及之后哪个消费者会为活动世代读取那个字段

活动世代通过完全独立的通道进入 compile-environment 设置:AcceleratorTypeToTpuVersionEnum @ 0x204cf620 把用户的 accelerator_type 字符串(v5ev5pv6etpu7x 等)解析为一个 TpuVersion ordinal;随后该 ordinal 驱动数据选择(通过 DefaultsForVersion @ 0x20b1b040 选择每世代 _chip_parts.binarypb 默认值)和 codec/HAL family 选择(TpuCodec::Create @ 0x1e835fa0),而不是 flag 选择。唯一真正按 TpuVersion 键控 flag 名称的位置,是一对很窄的 legacy-MSA 注册表(Legacy{Evictions,Prefetches}FlagRegistryflat_hash_map<TpuVersion, …>);其他所有内容都是 gen-blind。

对重实现而言,契约是:

  • 非门控模型 —— codename 前缀 flag 会无条件注册并无条件应用;前缀是命名空间,不是分发分支。不要构建按世代的注册 switch。
  • active-gen → codename → data 绑定 —— TpuVersion ordinal 如何降为 codename 字符串(TpuVersionToString),并馈入嵌入的 <codename>_chip_parts.binarypb 默认 proto 路径以及每 codename 的 codec/HAL family;这些都由同一张 ordinal 表键控。
  • 唯一的 gen-keyed 例外 —— legacy MSA eviction/prefetch 的每 TpuVersion flag 名称映射,以及为什么它们是二进制中唯一的 flat_hash_map<TpuVersion,…> flag 注册表。
Flag→field 桥(覆盖时应用)OverrideTpuCompEnvByCmdLineFlags @ 0x1d73e640
Flag→field 桥(构建默认值)CreateDefaultTpuCompEnv @ 0x1d73dfa0
逐字段 flag 查找(gen-blind)TpuCompEnvReflection::GetFlagForField @ 0x1d74ad40
逐 flag 字段查找(反向)TpuCompEnvReflection::GetFieldForFlag @ 0x1d74ab20
全局 flag↔field 映射FlagFieldMappings::GetInstanceabsl::NoDestructor,ctor @ 0x1d753ce0
逐 flag 写入器SetFieldFromFlagString @ 0x1d73fcc0
active-gen 入口(string→ordinal)AcceleratorTypeToTpuVersionEnum @ 0x204cf620libtpu_init_utils.cc
ordinal → codename 字符串TpuVersionToString @ 0x20b3a480 → 表 off_22011BF0
ordinal → 每 gen 默认 protoTpuChipParts::DefaultsForVersion @ 0x20b1b040
ordinal → codec/HAL familyTpuCodec::Create @ 0x1e835fa0
gen-keyed flag 名称映射(例外)LegacyEvictionsFlagRegistry @ 0x1c6f8760LegacyPrefetchesFlagRegistry @ 0x1c6f8940

速览:前缀、世代与门控

每前缀计数是来自 flag-families 的 registration-symbol-true 数字(AbslFlagHelpGenFor<prefix> 符号)。门控机制列是本页恢复出的内容:对每个路由到 TCE 的家族,答案都一样 —— 无条件注册,gen-blind 应用。codename 列记录的是作者关联,不是运行时过滤器。

前缀作者世代 / codename数量门控机制
xla_tpu_*无(所有 gen)909无条件注册;gen-blind 反射应用
xla_jf_*Jellyfish(v0)命名空间,all-gen 编译器核心148无条件注册;gen-blind 反射应用
xla_sc_*无(SparseCore 后端)92无条件注册;gen-blind 反射应用
barna_core_*无(embedding runtime)61无条件注册;gen-blind 反射应用
xla_vf_*Viperfish(v3)VMEM/MSA16无条件注册;gen-blind 反射应用
xla_gf_*6acc60406(TpuVersion 5,v7x)VMEM/MSA14无条件注册;gen-blind 反射应用
xla_pf_*Pufferfish(v2)ND all-reduce1无条件注册;gen-blind 反射应用
(legacy MSA evict/prefetch)TpuVersion flag-name 映射n/agen-keyed flat_hash_map<TpuVersion,…>
xla_gl_*(Ghostlite codename)0不存在 —— 无 flag,无字符串

注意 —— flag 注册或应用上没有 active-gen gating。没有按世代的注册 switch,也没有按世代的应用过滤器:所有 TCE flag,无论是否带 codename 前缀,都会无条件注册,并且只要出现在命令行上,就会被写入其 TCE 字段,不管活动 TpuVersion 是什么。xla_jf_* flag 在 Viperfish 上不会 no-op,也不携带 gen-specific 默认值。每世代行为完全存在于(a)按 codename 加载的默认 proto(DefaultsForVersion)以及(b)活动 gen 上各消费 pass 读取哪些字段。OverrideTpuCompEnvByCmdLineFlagsCreateDefaultTpuCompEnv 不接受 TpuVersion,反射循环遍历完整描述符。

陷阱 —— catalog 中的 xla_gl_* / Ghostlite codename 家族在此构建中不存在:零个 xla_gl_* 字符串,零个注册符号。codename ghostlite 作为硅片 codename 是 live 的(TpuVersion=4TpuVersionToString slot 4),但不保留 flag-prefix 命名空间。重实现者不得分配 gl flag 分支。live codename flag-prefix 集是 {jf, pf, vf, gf}


1. 分发机制

目的

经由 LIBTPU_INIT_ARGS 到达的 flag(例如 --xla_vf_max_vmem_used_by_memory_space_assignment=...)必须落入正确的 TpuCompilationEnvironment 字段。本节恢复这种绑定是如何执行的,并展示该绑定与世代无关。flag 名称中的 codename 只在作者命名时被消费(它消歧字段名);运行时 codename token 是不透明的,会在一个以 FieldDescriptor* 为键的映射中被哈希掉。

入口点

text
PJRT init / TpuPlatform setup
  └─ AcceleratorTypeToTpuVersionEnum (0x204cf620)   ── parse "v5e"/"tpu7x"/… → TpuVersion ordinal
  └─ CreateDefaultTpuCompEnv (0x1d73dfa0)           ── seed every TCE field from its flag's current value
       └─ TpuCompEnvReflection::GetFlagForField (0x1d74ad40)  ── field → flag, via FlagFieldMappings
       └─ TpuCompEnvReflection::ReadFlag                       ── flag value → variant
       └─ TpuCompEnvReflection::SetEnvField                    ── variant → TCE field
  └─ OverrideTpuCompEnvByCmdLineFlags (0x1d73e640)  ── re-walk; override fields whose flag WasPresentOnCommandLine
       └─ (same reflection trio)
       └─ SetFieldFromFlagString (0x1d73fcc0)        ── per-field string parse + write
```text

### 算法

覆盖桥是规范视图。注意循环边界,以及没有任何 `TpuVersion` 判别器。

```c
function OverrideTpuCompEnvByCmdLineFlags(env):        // sub_1d73e640
    meta = TpuCompilationEnvironment::GetMetadata()     // global descriptor metadata
    n    = meta.field_count                              // *(int*)(meta+8)
    if n <= 0: return OK
    overridden = []                                      // for the deprecation report
    for i in 0 .. n-1:                                   // walks EVERY TCE field; v8 += 88 per field
        field = meta.fields[i]                           // 88-byte FieldDescriptor record
        flag  = GetFlagForField(field)                   // sub_1d74ad40 — one map, no gen key
        if not flag.ok(): fail("Flag is not found for field")
        if not WasPresentOnCommandLine(flag):            // absl::flags_internal — user set it?
            continue                                     // untouched flags keep their default
        old = GetFieldValueAsString(field, env)
        // log "Overriding flag <name> to <new>; Old value was: <old>"
        if field.is_deprecated():                        // *(field+56)+125 == 1
            overridden.push(field.name)
        value = ReadFlag(flag, field)                    // current flag value → 20-arm variant
        SetEnvField(value, field, env)                   // write into env's proto field
    if overridden not empty:
        LOG(WARNING) << "[DEPRECATED_XLA_TPU_FLAG_USE] Deprecated "
                     << "TpuCompilationEnvironment flags were overridden: "
                     << join(overridden, ", ")
    return OK

CreateDefaultTpuCompEnv @ 0x1d73dfa0 是镜像:同样完整遍历描述符,但它从 ReadFlag(flag 的当前值;若未设置则为默认值)播种每个字段,而不是由 WasPresentOnCommandLine 门控。随后它通过 MessageDifferencer::CompareWithFields 将结果与 GetTpuCompEnvWithDefaultValues() diff,并发出并行警告 [DEPRECATED_XLA_TPU_FLAG_USE] Deprecated TpuCompilationEnvironment flags were present and not matching their default values:(调用点字符串,tpu_compilation_environment.cc 源码第 5776 行)。它唯一的非描述符参数是 SparseDenseMatmulFdoConfig*,没有 TpuVersion

注意 —— deprecation report 是二进制中最接近“此 flag 不适用于这里”信号的东西,但它按世代划分。某 flag 是否 deprecated 由其 FieldDescriptor 上的静态 bit 标记(offset +56 → +125),对每个 target 都相同。在“错误”世代上设置 codename flag 不会产生警告或错误:它会像在“正确”世代上一样被写入字段,只是之后无人读取。

Flag↔Field 映射

GetFlagForField @ 0x1d74ad40 通过一个首次使用时惰性构建的单进程全局结构,把 FieldDescriptor* 解析为其 absl::CommandLineFlag*

c
function GetFlagForField(field):                         // sub_1d74ad40
    once: FlagFieldMappings::GetInstance()               // absl::NoDestructor, ctor 0x1d753ce0
    map = FlagFieldMappings.singleton                    // Swiss-table: FieldDescriptor* → CommandLineFlag*
    h   = crc32_hash(field)                               // MixingHashState over the pointer
    slot = swiss_table_probe(map, field, h)              // SSE group scan (vpcmpeqb / vpmovmskb)
    if slot found: return flag
    return Error("Flag is not found for field")          // tpu_compilation_environment_reflection.cc:80
```text

`GetFieldForFlag @ 0x1d74ab20` 是反向(flag → field),供 `SetFieldFromFlagString` 使用。二者查询*同一个* `FlagFieldMappings` 实例。该映射以原始 `FieldDescriptor*` / `CommandLineFlag*` 身份为键,在 static-init 时构造一次,覆盖整个 TCE 表面;没有按世代的变体,也不会在 `TpuVersion` 改变时重建。这就是 codename 前缀不是运行时分发轴的结构性证明:当 flag 到达应用循环时,它的 `jf`/`vf`/`gf` token 已经折叠为一个指针键映射条目,该条目由所有世代共享。

### 函数图

| 函数 | 地址 | 角色 |
|---|---|---|
| `OverrideTpuCompEnvByCmdLineFlags` | `0x1d73e640` | 完整描述符遍历;覆盖已出现 flag;发出 deprecation report |
| `CreateDefaultTpuCompEnv` | `0x1d73dfa0` | 完整描述符遍历;从 flag 播种默认值;与规范默认值 diff |
| `SetFieldFromFlagString` | `0x1d73fcc0` | 每 flag:字段查找 → `ParseFlagFromString` → `SetEnvField` |
| `TpuCompEnvReflection::GetFlagForField` | `0x1d74ad40` | 通过 `FlagFieldMappings` Swiss-table 执行 field → flag |
| `TpuCompEnvReflection::GetFieldForFlag` | `0x1d74ab20` | flag → field(反向) |
| `FlagFieldMappings` ctor | `0x1d753ce0` | 构建全局 flag↔field 映射(一次) |

### 注意事项

贯穿每个反射调用的 20 臂 `std::variant`(`a/h/i/l/j/m/d/f/b/string/RangeSpecProto/RepeatedStrings/SparseDenseMatmulFdoConfig/SlicedPrefetchOptions/MemoryBoundLoopOptimizerOptions/PreferredPrefetchOverrides/MsaSortOrderOverrides/BufferContentsSanitizerConfig/BufferIsolationConfig/AutoProto`)是字段类型宇宙,不是世代宇宙。匹配二进制的重实现者必须复现 variant 类型列表和每类型的 `NormalizeFieldType<T>` 特化(每个臂一个 `.text` 函数,例如 `NormalizeFieldType<TristateFlag> @ 0x1d761080`),但其中任何一个都不需要 `TpuVersion` 条件逻辑。

---

## 2. Active-Gen → Codename 绑定

### 目的

活动世代是真实存在的,也确实驱动行为,只是不驱动 flag 选择。本节恢复从面向用户的 `accelerator_type` 字符串到 `TpuVersion` ordinal 的链条,以及从 ordinal 到它实际选择的两件事:每世代默认旋钮 proto 和每 codename 的 codec/HAL family。它们是按 ordinal 键控的*数据**工厂*选择,与 §1 的 flag 流水线并行但相互分离。

### 入口点 —— 字符串到 Ordinal

`AcceleratorTypeToTpuVersionEnum @ 0x204cf620`(位于 `libtpu::(anonymous)`,源码 `learning/45eac/tfrc/runtime/libtpu_init_utils.cc`)是世代第一次变成数字的位置。它按 `-` 拆分 `accelerator_type`(`"<version>-<cores>"`),将第一个 token 转为小写,并做映射:

```c
function AcceleratorTypeToTpuVersionEnum(accel_type):     // sub_204cf620
    parts = split(accel_type, '-')                         // "v5e-256" → ["v5e","256"]
    if parts.size != 2: return Error("...not in the format of '<tpu_version>-<core_count>'...")
    v = lower(parts[0])
    switch v:                                              // string compares, ordinals are TpuType-side
      "v4lite" -> 4     "v2"  -> 1     "v3"  -> 2     "v4"  -> 3
      "v5lite"/"v5e" -> 5     "v5p" -> 6
      "v6e"/"v6ea"   -> 7
      "tpu7x"/"tpu7" -> 8
      default        -> Error("Unsupported accelerator type: " + accel_type)
    return v

怪癖 —— 这个解析器发出的 ordinal 是 public superpod::routing::TpuType enum 值(v2→1,v5e→5,v5p→6,v6e→7,tpu7x→8),不是内部 TpuVersion enum(kJellyfish=0 … k6acc60406=5)。两个 enum 会在别处协调;重实现者必须保持它们不同。IsAtLeastTPU7x @ 0x204cfda0 是直接建立在这个解析器上的薄 capability gate:它返回 parsed_type >= 8,即 tpu7x/tpu7。这就是“是否至少是最新世代”检查的整个函数体。

Ordinal 到 Codename 字符串

TpuVersionToString @ 0x20b3a480 是带边界检查的表索引,是规范的 ordinal → codename 降低:

c
function TpuVersionToString(v):                  // sub_20b3a480
    if v >= 6: LOG(FATAL) << "Invalid TPU version " << v   // tpu_version.cc:152
    return off_22011BF0[v]                         // 6-entry .data.rel.ro pointer table
```text

六个 relocated 表项依次是:`jellyfish`、`dragonfish`、`pufferfish`、`viperfish`、`ghostlite`、`6acc60406`(`TpuVersion` 0..5)。这张表是把 ordinal 绑定到 codename 字符串的唯一真源;默认 proto 路径和诊断名称路径都通过它读取。

### Ordinal 到每 Gen 默认 Proto

`TpuChipParts::DefaultsForVersion @ 0x20b1b040` 是 codename 字符串变成*数据*选择器的位置。它最清楚地展示了每世代轴是数据,而不是 flag 路由:

```c
function DefaultsForVersion(version, variant):    // sub_20b1b040
    name = AsciiStrToLower(TpuVersionToString(version))    // ordinal → "viperfish"
    if variant: name = name + "_" + variant                // optional chip variant suffix
    path = "embed://tpu_chip_parts/" + name + "_chip_parts.binarypb"
    // e.g. "embed://tpu_chip_parts/viperfish_chip_parts.binarypb"
    proto = TpuChipPartsProto()
    ReadBinaryProto(Env::Default(), path, &proto)          // tpu_chip_parts.cc:343 (CHECK OK)
    return TpuChipParts::FromProto(proto)

嵌入资源名由降低后的 codename 组装而成:字面量 "embed://tpu_chip_parts/"(23 字节)和 "_chip_parts.binarypb"(20 字节)夹住 codename。因此每个世代都随二进制携带一个不同的默认参数 proto;选择世代就是选择哪个 proto 的默认值为编译播种。§1 中 CreateDefaultTpuCompEnv 建立的 TCE flag 默认值和这里的 chip-parts 默认值是彼此独立的默认值来源:flag 按其 absl::Flag<T> 静态默认值默认,chip parts 按每 gen 的 .binarypb 默认。

Ordinal 到 Codec / HAL Family

同一个 ordinal 还会选择 codec 和 HAL family,这是一个直接 switch,是 flag map 的 codec/HAL 对应物:

c
function TpuCodec::Create(version):               // sub_1e835fa0
    switch version:
      0: return CreateTpuCodecJellyfish()          // jxc family
      1: return CreateTpuCodecDragonfish()          // jxc family
      2: return CreateTpuCodecPufferfish()          // pxc family
      3: return CreateTpuCodecViperfish()           // vxc family
      4: return CreateTpuCodecGhostlite()           // gxc/glc family
      5: return sub_1E838380()                       // gxc/gfc family (6acc60406, anonymous class)
```text

这就是主题问题中提到的 codec/HAL-family 选择:它按 *ordinal* 键控,完全类似 chip-parts 路径,并且与 flag 分发不共享任何机制。codename → family-tag 关系(`jellyfish/dragonfish`→`jxc`,`pufferfish`→`pxc`,`viperfish`→`vxc`,`ghostlite/6acc60406`→`gxc`)与 flag 前缀是同源命名关系,但 family 选择发生在这里,不发生在 flag 层。

> **注意 ——** flag *前缀*与 codec/HAL *family*之间的关系纯属词法:`xla_vf_*` flag 和 `vxc`/Viperfish codec 都源于 `viperfish` codename,但二者互不选择。flag 前缀由写旋钮的工程师选择,用于把它限定到某个世代的调参空间;codec family 在运行时由 `TpuVersion` ordinal 选择。它们只在 codename 字符串处相遇,从不共享一个分发过程。

### 函数图

| 函数 | 地址 | 角色 |
|---|---|---|
| `AcceleratorTypeToTpuVersionEnum` | `0x204cf620` | 解析 `accelerator_type` 字符串 → `TpuType` ordinal |
| `IsAtLeastTPU7x` | `0x204cfda0` | capability gate:已解析 ordinal `>= 8` |
| `TpuVersionToString` | `0x20b3a480` | ordinal → codename 字符串(表 `off_22011BF0`) |
| `TpuChipParts::DefaultsForVersion` | `0x20b1b040` | ordinal → `<codename>_chip_parts.binarypb` 默认值 |
| `TpuCodec::Create` | `0x1e835fa0` | ordinal → 每 codename 的 codec / HAL family |

---

## 3. 唯一的 Gen-Keyed Flag 路径

### 目的

二进制中恰好有一个机制按 `TpuVersion` 键控 flag *名称*:一对 legacy memory-space-assignment 注册表。它们是 §1 的 gen-blind 规则的真正例外;恢复它们能避免重实现者漏掉 per-gen 间接层,或把它过度泛化到整个 flag 表面。

### 机制

`LegacyEvictionsFlagRegistry @ 0x1c6f8760` 和 `LegacyPrefetchesFlagRegistry @ 0x1c6f8940` 是 `util_registration::StaticMapBase` 特化,其值映射是:

```text
flat_hash_map<tpu::TpuVersion,
              std::pair<const char* /*flag name*/, string_view /*location*/>>

也就是 TpuVersion → flag-name 映射。在 static-init 时,每个拥有 legacy eviction(或 prefetch)旋钮的世代都会通过 InsertValue(TpuVersion, name, sourceloc)自己的 flag 名称按 ordinal 注册。该映射强制每个 key 只能定义一次;同一 TpuVersion 的重复注册会触发 LOG(FATAL),消息为 "Attempting to redefine value for key …"(static_map.h:141)。需要“活动世代的 legacy eviction policy flag”的消费者会在此映射中查找活动 TpuVersion,读取解析出的 flag,而不是硬编码某个 flag 名称。

为什么它是例外

这两个注册表存在,是因为 legacy MSA eviction/prefetch 旋钮早于统一的 TCE 反射表面,并保留了它们的每世代 flag 名称(而不是使用一个由条件消费的 TCE 字段)。它们是二进制中唯一的 flat_hash_map<TpuVersion, …> flag 注册表;主表面(909 个 xla_tpu_*、全部 codename 家族、xla_sc_*barna_core_*)都通过单个 gen-blind 的 FlagFieldMappings 路由。重实现者应按 §1 建模一般情况,并把这两个注册表视为边界明确的 legacy carve-out。

函数图

函数地址角色
LegacyEvictionsFlagRegistry::InsertValue0x1c6f8760TpuVersion eviction-flag-name 注册
LegacyPrefetchesFlagRegistry::InsertValue0x1c6f8940TpuVersion prefetch-flag-name 注册

陷阱 —— 不要泛化这些注册表。发现一个按 flat_hash_map<TpuVersion,…> 键控的 flag 注册表,可能会让人以为整个 flag 系统都是 gen-keyed;事实不是这样。这两个 sync_flag_util 映射只是 MSA eviction/prefetch 的局部 legacy 机制。其他所有 TCE flag,包括全部 jf/pf/vf/gf codename flag,都由 §1 中指针键控、generation-blind 的 FlagFieldMappings 解析。


相关组件

组件关系
FlagFieldMappings::GetInstance(ctor 0x1d753ce0唯一全局 flag↔field 映射;把 codename token 折叠成 FieldDescriptor*
TpuCompEnvReflection::ReadFlag / SetEnvFieldflag-value ↔ TCE-field-value 传递,逐字段,gen-blind
NormalizeFieldType<T> 家族(例如 0x1d761080每 variant 臂的类型强制转换;每字段类型一个,无每世代版本
GetTpuCompEnvWithDefaultValuesCreateDefaultTpuCompEnv diff 使用的规范 TCE 默认快照
off_22011BF0(TpuVersionToString 表)ordinal → codename 字符串;由默认 proto 和诊断路径共享

交叉引用

  • flag-families.md —— prefix → owner / sink 路由图;本页是其下方的分发层(注册和应用)
  • xla-flag-atlas.md —— 平坦逐 flag catalog,以及 codename 家族索引到的子系统关键词分类法
  • flag-catalog-full.md —— 建立本文引用的每前缀注册计数的 appendix
  • tpu-compilation-environment.md —— TCE 主 proto,本分发遍历其每个字段;字段布局和反射表面
  • registry-mediated-flags.md —— TpuCompEnvReflection flag↔field 桥,通过一条通用 gen-blind 路径服务所有 TCE 家族
  • overview.md —— 三层 config 流水线,以及 variant 臂携带的 AutoProto 三态旋钮
  • codename-cheatsheet.md —— TpuVersion ordinal ↔ codename ↔ public-name 表,是 TpuVersionToStringAcceleratorTypeToTpuVersionEnum 的背景