结构化稀疏槽位 (v5+)
地址适用于
libtpu-0.0.40-cp314wheel 中的libtpu.so(BuildID md589edbbe81c5b328a958fe628a9f2207d— 明确锚点;运行时报告的0.103无法在二进制中静态验证;未 strip — 保留完整 C++ 符号)。其他版本会不同。
摘要
本页覆盖 Viperfish 及以后版本 (v5+) 的 TensorCore 结构化稀疏数据通路:MXU 原生消费的 1:N(2:4 风格)权重稀疏,使稀疏 matmul 或 convolution 比其 dense 等价形式少执行 block_size× 个 systolic 步。该特性以 convolution/dot HLO 上携带的一个小属性进入 MXU — 一个 xla::SparsityConfig proto — 它会通过 shape inference,在 jellyfish(TPU TensorCore)lowering 中按严格约束集验证,最后把 block_size 打包因子穿入 matmul emitter 的 gain-latch 和 window-config 计算。不存在单独的“sparsity opcode”:稀疏性是 matmul 的一个属性,它会重塑 stationary operand,并在编译器标志控制下发射为 SME 结构化稀疏 outer-product 指令族。
此二进制中有三个子系统都会说“sparsity”,不得混淆。本页讲的是第一个。 (1) 这里的 TensorCore MXU 稀疏,由 xla::jellyfish::ConvolutionEmitter、SpatialMajorConvolution 和 MatrixMultiplyAccumulateFunctor 拥有。(2) SparseCore embedding 引擎(xla::tpu::sparse_core::CustomKernelEmitter,SCS/TAC/TEC sequencer)— 一个完全独立的 Part IX 子系统,其 "Sparsity only supported on 2nd minor dimension" 检查不是关于 MXU 的。(3) 一个无效的 NVGPU/NVVM 导入路径(mlir::nvgpu::MmaSparseSyncOp、mlir::NVVM::MmaSpOp、getSparsitySelector、sparsitySelector 属性),它因为 MLIR 库被整体链接而出现在二进制中,但目标是 NVIDIA mma.sp 硬件,永远不会从 TPU 后端到达。任务简介中询问的“selector operand”属于那个 NVIDIA 路径,而不是 TPU MXU 槽位;TPU MXU 没有逐步 selector operand — 它从 packed-kernel 布局本身推导保留 lane 的模式。这个区分是本页最重要的一点,并记录在 §5。
重新实现者必须遵守的结构化稀疏契约:
- 配置 proto —
xla::SparsityConfig携带可选的lhs和可选的rhs,每个都是嵌套的SparsityConfig_TensorSparsityConfig,含四个标量字段:num_non_zero、block_size、dimension、stride。由ParseSparsityConfig@0x1e4fb500解析,由SparsityConfigToString@0x1e5a50c0序列化。 - 1:N 限制 —
num_non_zero必须等于1;任何其他值都会在ShapeInference::InferConvolveShape@0x1e539040中以 "Only 1:N sparsity is currently supported." 拒绝。dense contraction extent 是block_size× stored extent。 - 约束关卡 —
ConvolutionEmitter::ValidateConvolutionWithSparseKernel@0x130d6300强制stride == 1、稀疏维度必须是布置在 sublanes 上的 kernel input-feature 维度、没有 spatial dims,并且feature_group_count == batch_group_count == 1。 - MXU 贯穿 —
block_size成为LatchKernelPossiblyPackedImpl@0x1312c2c0和SpatialMajorConvolution::UpdateWindowConfigAndMegacoreSplitDim@0x1316e380中的打包因子;二者都会CHECKinput-feature tiling 是block_size的倍数,且num_non_zero() == 1。 - 发射门控 — SME 结构化稀疏 outer-product 指令族由帮助字符串为
"Enable SME Structured sparsity outer product instructions."的标志门控,该字符串位于0xa00bed8。可用性为 v5+(Viperfish 起);v4 及更早版本没有 SME 路径。
| 子系统 | TensorCore MXU 结构化 (1:N) 稀疏 — 不是 SparseCore,不是 NVGPU |
| 配置载体 | xla::SparsityConfig proto(可选 lhs,可选 rhs) |
| 逐 tensor 字段 | num_non_zero, block_size, dimension, stride |
| 稀疏率 | 仅 1:N(num_non_zero == 1,N == block_size) |
| 首代 | v5+ / Viperfish (kViperfish=3) 起;Pufferfish (v4) 及更早版本没有 SME 路径 |
| 发射门控 | 标志 enable_sme_structured_sparsity_outer_product_instructions(help @ 0xa00bed8) |
| 所属 namespace | xla::jellyfish(TensorCore 后端) |
| 关键 validator | ConvolutionEmitter::ValidateConvolutionWithSparseKernel @ 0x130d6300 |
1. SparsityConfig proto 及其字段布局
结构化稀疏以附着在 convolution 上的 xla::SparsityConfig message 进入编译器(HLO convolution op 携带它;HloInstruction::sparsity_config() @ 0x1e5aa080 是 accessor,返回嵌入 proto 的引用)。该 message 恰好有两个可选 sub-message — 每个 matmul operand 一个:
message SparsityConfig {
optional TensorSparsityConfig lhs = 1; // presence bit 0 (mask |1)
optional TensorSparsityConfig rhs = 2; // presence bit 1 (mask |2)
message TensorSparsityConfig { // xla::SparsityConfig_TensorSparsityConfig
optional int64 num_non_zero = ?; // the "1" of 1:N
optional int64 block_size = ?; // the "N" of 1:N
optional int64 dimension = ?; // which logical dim is sparse
optional int64 stride = ?; // must be 1 (see §3)
}
}lhs / rhs 的存在性由 SparsityConfig 对象偏移 +0x10 处的 32-bit _has_bits_ word 跟踪(parser 对 lhs 写入 *((_DWORD*)cfg + 4) |= 1,对 rhs 写入 |= 2;见 §2)。指向嵌套 TensorSparsityConfig 对象的两个指针位于 qword 偏移 +3(lhs,+0x18)和 +4(rhs,+0x20);二者都通过 proto2::Arena::DefaultConstruct<SparsityConfig_TensorSparsityConfig> 惰性 arena 分配。
在每个 TensorSparsityConfig 内,四个 int64 字段和 message 自身的 _has_bits_ byte,根据 parser 的 store 和 ToString reader,布局如下:
| 字段 | 对象偏移 | _has_bits_(byte @ +0x10) | ToString vtable slot |
|---|---|---|---|
num_non_zero | +0x18 | bit 0 (|1) | vt[3] |
block_size | +0x20 | bit 1(合并 |3) | vt[4] |
dimension | +0x28 | bit 2 (|4) | vt[5] |
stride | +0x30 | bit 3 (|8) | vt[6] |
NOTE(bit 编号 / 偏移)— 本页所有 bit 位置都是 LSB-first:
_has_bits_的“bit 0”是0x1mask,“bit 1”是0x2,依此类推,与 parser 实际发出的|1/|2/|4/|8OR 相匹配。上面的逐字段对象偏移来自 protoClear/InternalSerialize/parser store(*(_QWORD*)(cfg + 24)、+ 32、+ 40、+ 48)。无法从二进制恢复 proto-runtime field numbers(proto3 lite 会丢弃 field-number → name map),所以 wire tags 未钉死;只有内存布局和文本格式是确定的。ToStringvtable slotsvt[3..6]是该 message 的 accessor thunk,这里仅用于确认字段身份。
文本形式
SparsityConfigToString(@ 0x1e5a50c0)是规范 serializer,也是字段语义最清晰的单一证据。它对每个存在的 operand 发出:
lhs={sparsity=<num_non_zero>x<block_size> dimension=<dimension> stride=<stride>}
rhs={sparsity=<num_non_zero>x<block_size> dimension=<dimension> stride=<stride>}当二者都存在时用单个空格连接。sparsity= 值字面上是 FastIntToBuffer(num_non_zero) + "x" + FastIntToBuffer(block_size)(0x1e5a50c0:95-111),所以打印为 sparsity=1x4 的配置意味着 num_non_zero=1, block_size=4 — 即每四个一组保留一个非零值,经典 1:4 / 2:8 等价结构化模式。dimension= 读取 vt[5],stride= 读取 vt[6]。
2. ParseSparsityConfig
HLO 文本 parser 入口点是 HloParserImpl::ParseSparsityConfig @ 0x1e4fb500(file/line 字符串将其钉在 hlo_parser.cc:7164;嵌套的逐 operand parse 记录为 ParseTensorSparsityConfig,位于 :7199)。它实现语法 sparsity_config={lhs={…}, rhs={…}}:
// ParseSparsityConfig(SparsityConfig* out) @ 0x1e4fb500
expect '{'; // open brace, else return false
while (peek != '}') {
name = lex_identifier(); // must be a token, else "expects attribute name"
if (name == "lhs") { out->_has_bits_ |= 1; cfg = out->mutable_lhs(); }
else if (name == "rhs") { out->_has_bits_ |= 2; cfg = out->mutable_rhs(); }
else { Error("unknown attribute"); return false; }
expect '{'; // open the TensorSparsityConfig body
while (peek != '}') { // ParseTensorSparsityConfig body
attr = lex_identifier();
if (attr == "stride") { cfg->set_stride( ParseInt64() ); cfg->_has |= 8; }
else if (attr == "dimension") { cfg->set_dimension( ParseInt64() ); cfg->_has |= 4; }
else if (attr == "sparsity") { // value is "num_non_zero:block_size"
auto [nnz, blk] = ParseDxD(); // expects 'num_non_zero:block_size'
cfg->set_num_non_zero(nnz); cfg->_has |= 1;
cfg->set_block_size(blk); cfg->_has |= 3;
}
else { Error("unknown attribute"); return false; }
}
expect '}';
}
expect '}';有三个细节值得重新实现者注意。第一,operand key 通过针对 "lhs"(0x686C,'s')和 "rhs"(0x6872,'s')的内联 3-byte 比较匹配 — 恰好两个,没有 out 或 both。第二,sparsity 值由 ParseDxD 解析,错误文本是 "expects 'num_non_zero:block_size'"(0x1e4fb500:329),确认冒号分隔的 N:M 语法,其中第一个数是 num_non_zero,第二个数是 block_size。第三,内部属性比较对象是 "stride"(0x69727473+0x6564)、"dimension"(0x6F69736E656D6964+'n')和 "sparsity"(0x7974697372617073)— 只有这三个;num_non_zero 和 block_size 不是独立的文本属性,它们是 sparsity= 值的两半。
GOTCHA(selector)— TPU 语法中没有
selector属性。sparsitySelector属性字符串以及 "Invalid attributesparsitySelectorin property conversion:" 错误(@0xa255171)属于mlir::nvgpu::MmaSparseSyncOp::setPropertiesFromAttr@0x17067b40,即无效 NVGPU 路径(§5)。TPU sparse matmul 不携带逐指令 selector operand;保留 lane 的 mask 隐含在 packed kernel 如何存储于 sublanes 上(§4)。
3. Shape inference 和约束关卡 {#3-shape-inference-and-the-constraint-gauntlet}
Shape inference:仅 1:N,contraction × block_size
ShapeInference::InferConvolveShape @ 0x1e539040(来自 shape_inference.cc 的日志)是计算 sparse convolution 输出形状并施加第一个硬限制的位置。到达 RHS-with-sparsity 分支后,它读取 num_non_zero(vtable slot vt[3])并要求其为 1:
// InferConvolveShape, sparse-RHS branch @ 0x1e539040:1262
auto* tsc = sparsity_config.rhs(); // TensorSparsityConfig*
if (tsc->num_non_zero() == 1) // vt[3] == 1
output_feature_extent *= tsc->block_size(); // vt[4]; dense extent = stored × N
else
return InvalidArgument("Only 1:N sparsity is currently supported."); // @0x1e539040:1271语义是精确的:stored(compressed)RHS 每个 block_size 组保存一个保留值,所以 dense contracting extent 是 block_size × stored_extent。重新实现者必须在推断 matmul 的 K 维度时应用这个乘法,否则下游 tiling 计算会偏离 block_size。
QUIRK(1:N,不是 2:4)— 尽管“2:4-style”是口语名称,二进制支持的是更通用的 1:N 形式,而不是硬编码 2:4。
num_non_zero被钉为1,但block_size(N)是自由的int64(sparsity=1x4、1x8、…)。下游 MXU emitter 还会CHECKKernelSparsityConfig().num_non_zero() == 1(SpatialMajorConvolution::UpdateWindowConfigAndMegacoreSplitDim@0x1316e380:2980),所以 1:N invariant 被断言两次 — 一次是 shape inference 中面向用户的错误,一次是 lowering 中的内部CHECK。
jellyfish validator
shape inference 通过后,TensorCore lowering 会在 ConvolutionEmitter::ValidateConvolutionWithSparseKernel @ 0x130d6300(convolution_emitter.h)中用远更严格的集合重新验证。下面每条约束都是单独的 InvalidArgument,有自己的 format string;全部由这个函数拥有:
| 约束 | 拒绝消息(节略) |
|---|---|
| Sparse dimension == kernel input-feature dim | "expected kernel input feature dimension to be the sparse dimension." |
| Sparse dim 布置在 sublanes 上 | "expected kernel sparse dimension to be on sublanes." |
stride == 1 | "sparse conv with kernel_sparsity_stride not equal to 1 is not supported yet." |
| 无 spatial dims | "sparse conv with spatial dimensions is not supported yet." |
feature_group_count == batch_group_count == 1 | "…feature_group_count or batch_group_count not equal to 1 is not supported yet." |
| batch 是 N 的倍数 | "expected batch to be a multiple of %d." |
| input-feature 是 N 的倍数 | "expected input feature to be a multiple of %d." |
| kernel / indices type 在允许集合内 | "expected kernel type to be one of %s and indices type to be one of %s." |
合起来看,这些约束说明:此 TPU 上的 structured-sparse matmul 是一个 pointwise(1×1,无 spatial window)convolution,其 RHS(kernel)沿 input-feature 轴压缩,该轴位于 VMEM tile 的 sublane 维度上,compression stride 为 1,没有 feature/batch grouping,并且 batch 和 input-feature extent 都是 block factor N 的整数倍。batch/input-feature 消息中的 %d 倍数就是 block_size。
NOTE(kernel-as-tuple)— packed sparse kernel 作为 2-tuple operand 到达:
{values, indices}(compressed weights 加上每个 block 的 kept-position metadata)。TpuInstructionFusion::BitcastConvOperands以 "Got tuple as kernel but no kernel sparsity config"(字符串 @0x86645ed)保护这一点,shape-inference RHS 路径也以 "rhs of convolution, if a tuple, must have 2 elements for sparsity" 保护。因此 sparsity metadata 是真实的第二个 operand,而不是 side attribute — 但它以结构方式消费,不是每次乘法的数值 selector。
4. MXU 贯穿:block_size 打包因子 {#4-mxu-threading-the-block_size-packing-factor}
在 matmul emitter 内,block_size 是重塑 systolic schedule 的唯一数字。两个函数携带它。
SpatialMajorConvolution::UpdateWindowConfigAndMegacoreSplitDim @ 0x1316e380 在调整 window/megacore split 之前会 CHECK KernelSparsityConfig().num_non_zero() == 1(:2980),使 compressed input-feature 轴按 block_size 大小的 chunk 行走,而不是按 dense extent。所谓“packed”表示 MXU 的 stationary operand pool 只持有 dense weights 的 1/block_size,所以一个 latch + matmul 序列覆盖的 dense contraction 比 dense 形式多 block_size×。
MatrixMultiplyAccumulateFunctor::LatchKernelPossiblyPackedImpl @ 0x1312c2c0(名称中的 PossiblyPacked 就是 sparsity hook)是 gain-latch builder,与 dense 路径共享。它的两个 sparsity-specific assertion 钉住 tiling 契约:
// LatchKernelPossiblyPackedImpl @ 0x1312c2c0
CHECK(effective_input_feature_sublane_chunks_per_tile
% conv_->KernelSparsityConfig().block_size() == 0); // :129
CHECK(effective_input_feature_tile_start_in_chunks
% conv_->KernelSparsityConfig().block_size() == 0); // :151二者都要求 input-feature sublane chunking — 每 tile 的数量 和 tile start offset — 精确为 block_size 的倍数。这就是 §3 中 "input feature to be a multiple of %d" 约束的物理原因:packed kernel 以一个 block_size-chunk 为单位加载进阵列,partial chunk 没有有效 latch encoding。随后 emitter 发出 SME 结构化稀疏 outer-product 指令族,而不是 dense MXU matmul;保留 lane 模式由 packed layout(sublanes 上的 values+indices)携带,所以不会向 bundle slot 本身添加额外 operand field。
NOTE — 没有新槽位。 虽然此主题被 frontier register 命名为“slot”,但在 bundle 层面,结构化稀疏不会新增 TensorCore bundle slot,也不会新增 operand field。它复用同一个 MXU
VectorExtendedslot(MXU slot)和同一组 gain-latch sub-slots(matprep / IAR / latch);不同之处在于 (a) 选择的指令族(SME outer-product vs dense matmul,由 emission flag 选择)以及 (b) 编织进 tiling/latch 计算的block_size除数。v5+ 64-bit bundle 中没有“sparsity field” — 使 matmul 变稀疏的状态存在于 operand 的 packed memory layout(sublanes 上的 values+indices),而不是某个 bundle bit。
5. 什么不是 TPU MXU 路径:NVGPU 和 SparseCore {#5-what-is-not-the-tpu-mxu-path-nvgpu-and-sparsecore}
有三组字符串看起来相关,但其实不相关。钉住它们是恢复工作的一部分。
NVGPU / NVVM (mma.sp) — 无效 NVIDIA 路径
| 符号 | 所属方 | 为什么它不是 TPU 路径 |
|---|---|---|
getSparsitySelector | mlir::nvgpu::MmaSparseSyncOp @ 0x17052500 | 面向 NVIDIA mma.sync sparse 的 NVGPU dialect op;从不会在 TPU 上 lowering |
sparsitySelector attr | MmaSparseSyncOp::setPropertiesFromAttr @ 0x17067b40 | "Invalid attribute sparsitySelector in property conversion:" 错误 @ 0xa255171 是 NVGPU property conversion |
| "sparsity selector should be 0 or 1" | MmaSparseSyncOp::verify @ … | NVIDIA 的 2-bit metadata selector,不是 TPU 概念 |
| "sparsity selector must be i32 type" | mlir::NVVM::MmaSpOp::verify @ 0x1658c500 | mma.sp 的 NVVM intrinsic verifier;目标是 PTX |
这些内容之所以随包发布,是因为 libtpu 静态链接了上游 MLIR NVGPU/NVVM dialect 库;其中没有任何内容能从 jellyfish TPU 后端到达。原始主题框架中的“selector operand”就是这个 selector — NVIDIA tensor core 读取的每 pair 2 bit metadata index — 它没有 TPU 对应物。TPU 把保留 lane 信息保存在 packed operand layout 中。
SparseCore — 完全不同的子系统
xla::tpu::sparse_core::CustomKernelEmitter::ChooseWindowLayout 拥有 "Sparsity only supported on 2nd minor dimension"(字符串 @ 0x85fbe43)。那是 SparseCore embedding/scatter-gather 引擎(SCS/TAC/TEC sequencer),一个独立的 Part IX 子系统,有自己的 custom-kernel emitter。它的“sparsity”是 sparse embeddings,不是 MXU structured weight sparsity,其 2nd-minor-dimension 规则与上面的 block_size/sublane 规则无关。不要合并二者。
6. 按代际的可用性
可用性信号是 SME 结构化稀疏 outer-product 指令族及其 emission flag。
| External name | Codename | TpuVersion ordinal | SME structured sparsity |
|---|---|---|---|
| TPU v2 / v3 | Jellyfish / Dragonfish | kJellyfish=0 / kDragonfish=1 | 否 |
| TPU v4 | Pufferfish | kPufferfish=2 | 否 |
| TPU v5 | Viperfish | kViperfish=3 | 是 |
| TPU v6 lite | Ghostlite | kGhostlite=4 | 是 |
| TPU7x | 6acc60406 | k6acc60406=5 | 是 |
门控本身是一个编译器标志,其帮助字符串 "Enable SME Structured sparsity outer product instructions." 位于 0xa00bed8(identifier aEnableSmeStructuredSparsityOuterProductInstructions;字符串偏移已在二进制中按字节确认)。该帮助字符串只从 flag-registration table 到达,而不是从任何反编译函数到达 — 扫描反编译语料未找到对该字符串或 enable_sme_structured_sparsity_outer_product_instructions identifier 的任何 code reference — 确认它是一个boolean compiler flag,切换 lowering 是否可以发射 SME family,而不是逐调用 runtime branch。
NOTE(gate vs gen — 边界为推断)— 该标志是 emission gate;hardware gate 隐含在哪些代际有 SME unit(v5+)中。在没有 SME unit 的代际上,该标志没有可发射的指令。此轮没有把最早报告 SME 支持的确切
TpuVersionenum ordinal 追踪到单个比较点;v5+/Viperfish 边界是基于 SME family 是 v5 时代新增项以及 v3/v4 MXU encoder 可达路径中没有任何 SME sparsity reference 的推断。与此一致,VF 64-bit Bundle 未在 ViperfishVectorExtended0families 中记录Sparsityop,并把 VF-specific backing finding 明确留作开放项。两页在同一事实上一致:这里记录的 source-level SME/jellyfish机制是 generation-agnostic,而逐代际可用性(哪些 silicon 实际带有 SME unit)是较软、未确认的部分。GOTCHA(block_size 默认值)— 该标志自身没有数值
block_size。block_size是逐 operand 的,由SparsityConfig提供,只受num_non_zero == 1、batch/input-feature 是N的倍数规则,以及 §4 中 sublane-chunk 可整除性的约束。常见的1x4是约定,不是硬件常量;二进制没有钉住固定N。
7. 重新实现清单
| 步骤 | 要做什么 | 锚点 |
|---|---|---|
| 1. 携带配置 | 将 SparsityConfig{ rhs={num_non_zero=1, block_size=N, dimension=ifeat, stride=1} } 附着到 conv | ParseSparsityConfig @ 0x1e4fb500 |
| 2. 推断形状 | 将 dense contraction extent 乘以 block_size;拒绝 num_non_zero != 1 | InferConvolveShape @ 0x1e539040 |
| 3. 验证 | 强制:pointwise(无 spatial)、sparse dim = sublanes 上的 kernel-ifeat、stride==1、无 groups、batch 与 ifeat 是 N 的倍数 | ValidateConvolutionWithSparseKernel @ 0x130d6300 |
| 4. 提供 metadata | 将 kernel 作为 {values, indices} 2-tuple 传入;values 是 compressed weights | BitcastConvOperands(tuple guard @ 0x86645ed) |
5. 贯穿 block_size | 以 block_size sublane chunks 行走 input-feature 轴;断言可整除 | LatchKernelPossiblyPackedImpl @ 0x1312c2c0, UpdateWindowConfig… @ 0x1316e380 |
| 6. 发射 | 在现有 MXU slot 上发出 SME structured-sparsity outer-product family,由该标志门控 | help @ 0xa00bed8 |
交叉引用
- MXU 槽位 — sparse 路径复用的
VectorExtendedmatmul slot;没有向其添加 sparsity field。 - Matprep、IAR 和 Latch 子槽位 — gain-latch 机制,其 tiling 计算采用
block_size除数。 - VF 64-bit Bundle — Viperfish (TPU v5) bundle format;首个带 SME 路径的代际。没有逐 bundle sparsity field。
- GL Bundle / GF Bundle — Ghostlite (TPU v6 lite) / 6acc60406 (TPU7x) bundle formats;共享 SME 路径的 v5+ family 成员。
- TpuVersion / Codename Matrix — 权威 enum↔codename↔external-name 映射;说明为什么
Trillium/Ironwood在二进制中出现次数为零。 - Instruction-Bits Master Database — 逐代际 field registry;确认不存在专用 sparsity slot。
- ISA Overview — slot taxonomy 以及 v3 vs v4/v5+ encoder lineage 分叉。
- Dot/Conv → MXU Lowering — 此 sparsity 属性穿过的 matmul-lowering pipeline。
- RaggedDot & Convolution Lowering — 拥有
ValidateConvolutionWithSparseKernel的 convolution emitter family。