软件栈XLA
XLA
也称为 XLA、Accelerated Linear Algebra、StableHLO、XLA 编译器
先用大白话
XLA 是把模型计算图编译成 TPU 机器码的编译器，决定每个操作如何切分、补齐、融合和调度。
技术定义
XLA 是面向线性代数的领域专用编译器，接收 StableHLO 图并为 TPU 和其他加速器生成可执行代码，负责融合、布局、切分和调度。
工程细节
每条 TPU 服务路径最终都落到 XLA。在 JAX 路径中，jax.jit 捕获一步计算后交给 XLA；在 TorchTPU 中，TorchDynamo 和 AOTAutograd 生成 FX 图，下沉为标准化中间表示 StableHLO，再由 XLA 编译。两种情况下编译器都是 XLA，不涉及 Inductor 和 Triton。编译器会把小于 MXU tile 的维度补齐、选择默认布局并融合逐元素操作，这对规则形状很高效，对与 tile 几何冲突的形状则代价高昂。Pallas 就是为 XLA 默认行为不够好的场景准备的。
为什么重要
编译器与硬件的协同设计是 Google 每 token 成本优势的重要来源：芯片、网络和编译器一起设计，计算和通信可以作为一个系统来优化。同样的紧耦合也意味着模型形状和编译 shape bucket 对性能有超出寻常的影响。
如何在 InferenceX 中解读
Official Preview 中的多项优化是在引导 XLA 而不是替代它，例如把专家 ID 和 token 索引打包成一个排序键，让 XLA 执行更简单的排序，排序延迟从 106.6 微秒降到 21.7 微秒，同时启用了 FP8 all-gather。