Skip to content

结构化稀疏槽位 (v5+)

地址适用于 libtpu-0.0.40-cp314 wheel 中的 libtpu.so(BuildID md5 89edbbe81c5b328a958fe628a9f2207d — 明确锚点;运行时报告的 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::ConvolutionEmitterSpatialMajorConvolutionMatrixMultiplyAccumulateFunctor 拥有。(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::MmaSparseSyncOpmlir::NVVM::MmaSpOpgetSparsitySelectorsparsitySelector 属性),它因为 MLIR 库被整体链接而出现在二进制中,但目标是 NVIDIA mma.sp 硬件,永远不会从 TPU 后端到达。任务简介中询问的“selector operand”属于那个 NVIDIA 路径,而不是 TPU MXU 槽位;TPU MXU 没有逐步 selector operand — 它从 packed-kernel 布局本身推导保留 lane 的模式。这个区分是本页最重要的一点,并记录在 §5

重新实现者必须遵守的结构化稀疏契约:

  • 配置 protoxla::SparsityConfig 携带可选的 lhs 和可选的 rhs,每个都是嵌套的 SparsityConfig_TensorSparsityConfig,含四个标量字段:num_non_zeroblock_sizedimensionstride。由 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 @ 0x1312c2c0SpatialMajorConvolution::UpdateWindowConfigAndMegacoreSplitDim @ 0x1316e380 中的打包因子;二者都会 CHECK input-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 == 1N == block_size
首代v5+ / Viperfish (kViperfish=3) 起;Pufferfish (v4) 及更早版本没有 SME 路径
发射门控标志 enable_sme_structured_sparsity_outer_product_instructions(help @ 0xa00bed8
所属 namespacexla::jellyfish(TensorCore 后端)
关键 validatorConvolutionEmitter::ValidateConvolutionWithSparseKernel @ 0x130d6300

1. SparsityConfig proto 及其字段布局

结构化稀疏以附着在 convolution 上的 xla::SparsityConfig message 进入编译器(HLO convolution op 携带它;HloInstruction::sparsity_config() @ 0x1e5aa080 是 accessor,返回嵌入 proto 的引用)。该 message 恰好有两个可选 sub-message — 每个 matmul operand 一个:

text
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 偏移 +3lhs+0x18)和 +4rhs+0x20);二者都通过 proto2::Arena::DefaultConstruct<SparsityConfig_TensorSparsityConfig> 惰性 arena 分配。

在每个 TensorSparsityConfig 内,四个 int64 字段和 message 自身的 _has_bits_ byte,根据 parser 的 store 和 ToString reader,布局如下:

字段对象偏移_has_bits_(byte @ +0x10ToString vtable slot
num_non_zero+0x18bit 0 (|1)vt[3]
block_size+0x20bit 1(合并 |3vt[4]
dimension+0x28bit 2 (|4)vt[5]
stride+0x30bit 3 (|8)vt[6]

NOTE(bit 编号 / 偏移)— 本页所有 bit 位置都是 LSB-first_has_bits_ 的“bit 0”是 0x1 mask,“bit 1”是 0x2,依此类推,与 parser 实际发出的 |1/|2/|4/|8 OR 相匹配。上面的逐字段对象偏移来自 proto Clear/InternalSerialize/parser store(*(_QWORD*)(cfg + 24)+ 32+ 40+ 48)。无法从二进制恢复 proto-runtime field numbers(proto3 lite 会丢弃 field-number → name map),所以 wire tags 未钉死;只有内存布局和文本格式是确定的。ToString vtable slots vt[3..6] 是该 message 的 accessor thunk,这里仅用于确认字段身份。

文本形式

SparsityConfigToString(@ 0x1e5a50c0)是规范 serializer,也是字段语义最清晰的单一证据。它对每个存在的 operand 发出:

text
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={…}}

c
// 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 比较匹配 — 恰好两个,没有 outboth。第二,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_zeroblock_size 不是独立的文本属性,它们是 sparsity= 值的两半。

GOTCHA(selector)— TPU 语法中没有 selector 属性。sparsitySelector 属性字符串以及 "Invalid attribute sparsitySelector in 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

c
// 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_sizeN)是自由的 int64sparsity=1x41x8、…)。下游 MXU emitter 还会 CHECK KernelSparsityConfig().num_non_zero() == 1SpatialMajorConvolution::UpdateWindowConfigAndMegacoreSplitDim @ 0x1316e380:2980),所以 1:N invariant 被断言两次 — 一次是 shape inference 中面向用户的错误,一次是 lowering 中的内部 CHECK

jellyfish validator

shape inference 通过后,TensorCore lowering 会在 ConvolutionEmitter::ValidateConvolutionWithSparseKernel @ 0x130d6300convolution_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 stride1,没有 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 契约:

c
// 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 VectorExtended slot(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 路径
getSparsitySelectormlir::nvgpu::MmaSparseSyncOp @ 0x17052500面向 NVIDIA mma.sync sparse 的 NVGPU dialect op;从不会在 TPU 上 lowering
sparsitySelector attrMmaSparseSyncOp::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 @ 0x1658c500mma.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 nameCodenameTpuVersion ordinalSME structured sparsity
TPU v2 / v3Jellyfish / DragonfishkJellyfish=0 / kDragonfish=1
TPU v4PufferfishkPufferfish=2
TPU v5ViperfishkViperfish=3
TPU v6 liteGhostlitekGhostlite=4
TPU7x6acc60406k6acc60406=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 支持的确切 TpuVersion enum ordinal 追踪到单个比较点;v5+/Viperfish 边界是基于 SME family 是 v5 时代新增项以及 v3/v4 MXU encoder 可达路径中没有任何 SME sparsity reference 的推断。与此一致,VF 64-bit Bundle 未在 Viperfish VectorExtended0 families 中记录 Sparsity op,并把 VF-specific backing finding 明确留作开放项。两页在同一事实上一致:这里记录的 source-level SME/jellyfish 机制是 generation-agnostic,而逐代际可用性(哪些 silicon 实际带有 SME unit)是较软、未确认的部分。

GOTCHA(block_size 默认值)— 该标志自身没有数值 block_sizeblock_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} } 附着到 convParseSparsityConfig @ 0x1e4fb500
2. 推断形状将 dense contraction extent 乘以 block_size;拒绝 num_non_zero != 1InferConvolveShape @ 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 weightsBitcastConvOperands(tuple guard @ 0x86645ed
5. 贯穿 block_sizeblock_size sublane chunks 行走 input-feature 轴;断言可整除LatchKernelPossiblyPackedImpl @ 0x1312c2c0, UpdateWindowConfig… @ 0x1316e380
6. 发射现有 MXU slot 上发出 SME structured-sparsity outer-product family,由该标志门控help @ 0xa00bed8

交叉引用