当我们提到常见的 GPU 运算——矩阵乘法、注意力机制、物理模拟——哈希通常不在第一反应里。哈希算法涉及复杂的计算和繁琐的数据依赖，很难做到并行化。虽然 GPU 已经被用于哈希运算（比如比特币 GPU 矿机），但并发通常只停留在"多个哈希同时算"的层面，也就是同时计算多个哈希，而不是针对单个摘要进行深度优化。
然而 Pearl 挖矿让我们别无选择：矿工需要提交的张量通常很大，而且已经驻留在显存里。如果在 CPU 端做哈希，每次矩阵乘法都会带来大量的 CPU-GPU I/O 开销。因此，要在 Pearl 挖矿中实现合理的开销，需要满足三个条件：
- 一个 GPU kernel，能够计算密码学安全的哈希
- 在单个输入上最大限度地发挥 GPU 的能力
- 不产生主机 I/O
本文详细介绍了我们在开发过程中做出的设计选择——从选择哈希算法、架构设计，到使用哪些 PTX 指令进行数据传输。
关于作者：Yoray Herzberg（@sag0li）是 @prlnet 核心团队的 kernel 工程师，目前在攻读 AI 安全硕士学位，在 Vaktibabat 上撰写密码学、系统和 GPU 编程相关内容。
BLAKE3
Pearl 选择使用 BLAKE3 哈希算法，因为它相比 SHA256 等哈希有极大的速度优势，而且更容易并行化。
BLAKE3 是一种基于 Merkle Tree 的哈希：它将数据分成 1024 字节的块，每块压缩后产生 32 字节的链值（chaining value），最后在这些链值上构建 Merkle Tree，迭代哈希每对节点，直到得到最终哈希。
每个 1024 字节的块被分成 64 字节的子块进行压缩。每个子块经过一个由 7 轮运算组成的压缩函数（包含 XOR、位移位等操作，轮间有置换）。除了 64 字节的消息外，压缩函数还需要一个 32 字节的链值（可以理解为块压缩过程中的"状态"），以及若干其他参数，如标志集、块大小和每个块独有的计数器。
在压缩一个块时，初始链值默认设为某个 IV，或在 keyed hashing 模式下设为用户自定义的 key，然后压缩第一个子块。接着，将输出链值作为下一个子块的输入链值，直到完成整个块的压缩。在 Pearl 中，我们始终使用 keyed hashing 模式，key 设为区块链中前一个块的哈希，以使噪声依赖于前一个块。
前一步（我们称之为叶子压缩）的最终结果是：对输入的每个 1024 字节块，产出一个 32 字节的输出链值。然后在 Merkle Tree 阶段，将每两个连续的链值拼接成一个 64 字节的块，使用相同的压缩函数但不同的标志进行压缩。重复此过程直到只剩一个链值，这就是数据的最终 BLAKE3 哈希。
树的归约过程取决于叶子节点的数量，但根据 BLAKE3 规范，必须满足以下两条规则：
- 左子树是满的：每个左子树都是完全二叉树，所有块位于同一深度，块数量为 2 的幂。
- 左子树是大的：每个左子树包含的块数大于等于其右兄弟子树的块数。
我们的实现
我们将实现分为 3 个 kernel：
- MerkleTreeRootsKernel：这是承担大部分工作的 kernel。除了压缩所有块（Merkle Tree 的叶子节点）——这是一个"令人尴尬地并行"的操作——它还融合了部分 Merkle Tree 归约，以避免不必要的全局内存往返。
- ComputeBlakeMTKernel：接收上一步输出的子树部分根值，执行额外的归约。
- ReduceRootsKernel：在上一步之后仍有多个子树的情况下执行最终归约，输出最终的 32 字节摘要。
压缩函数
如前所述，块压缩是一个令人尴尬地并行的操作：各块完全独立，因此可以让每个线程压缩一个块。为了优化压缩，每个线程完全在寄存器中完成压缩，压缩逻辑通过宏完全展开。每次压缩消耗 48 个寄存器用于状态（rState）、块（rBlock）和置换临时区（rOrigBlock），均为 CuTe 寄存器张量。7 轮 × 8 个 G 函数 × 每个 6 个操作 = 每个 64 字节子块 336 条 ALU 指令。
一个块 = 1024 字节 = 16 个子块。因此每个线程执行 16 次顺序压缩，向前传播链值。第一个子块带有 CHUNK_START 标志，最后一个带有 CHUNK_END。初始链值是区块链 tip 的哈希，存储在只读常量内存中以实现最快访问。
Kernel 1：MerkleTreeRootsKernel
这是主 kernel，几乎所有时间都花在这里。每个 CTA 的"任务"是压缩 2^t 个块（t 是可调的超参数），然后用 Merkle Tree 归约它们。t 同时控制每个 CTA 的线程数：128 + 2^t，其中 128 个生产者线程负责通过 TMA 从 GMEM 加载数据到 SMEM，2^t 个消费者线程各压缩一个块。
TMA 加载
Hopper 架构的 Tensor Memory Accelerator（TMA）在不占用任何计算资源的情况下处理全局内存到共享内存的传输。生产者线程在 kernel 启动时配置 TMA 描述符。
全局内存被看作一个 (num_chunks, 256) 的 uint32_t 矩阵，即每行是一个 1024 字节的块。每次 TMA 加载取一个 (2^t, kNumWordsPerLoad) 的 tile。例如，128 字节的加载大小下，每次迭代为每个消费者加载 32 个字，1024/128=8 次迭代完成所有块的加载。
消费者循环
每个消费者线程运行一个类似的循环：等待数据、从共享内存读取、压缩、释放屏障以通知生产者可以加载下一批数据。消费者在释放前进行同步，确保所有消费者在任何线程通知生产者之前已经完成当前阶段的读取。
在每个流水线阶段内，消费者使用 128 位加载将其共享内存切片读入寄存器，然后运行 BLAKE3 压缩。
共享内存布局与 Swizzle
共享内存存储两部分内容：
- smem_a：消费者压缩的数据。布局为 [kTmaThreads][kNumWordsPerLoad][kPipelineStages]，带 swizzle。每个消费者拥有一行（当前加载的块数据），第三维是流水线阶段。
- smem_leaves：部分 Merkle Tree 阶段的工作区。布局为 [8][kNumConsumerThreads]，带 SW128 swizzle，每个线程存储 8 个链值 uint32_t。
我们根据加载大小选择 swizzle 原子，以最小化 bank 冲突。
双流水线
为了绕过 TMA 大小的硬件限制，支持超过 256 个消费者，我们将消费者分成两组独立的 256 个，每组有自己的共享内存区域、流水线实例和 GMEM 视图。
生产者 warp 在同一次循环迭代中交错加载两组的数据。消费者组通过独立的命名屏障（PrimaryConsumers 和 SecondaryConsumers）分别同步。这只影响块压缩阶段；该阶段结束后，所有消费者线程照常执行 Merkle Tree 归约。
不完整块的处理
之前忽略了一个问题：网格中最后一个 CTA 处理的块数可能少于 2^t（当总块数不能被 2^t 整除时）。更进一步，它处理的最后一个块甚至可能是不完整的（<1024 字节），甚至不到一个 64 字节的子块。为了解决这个问题，分配到该块的消费者线程在每次 TMA 加载完成后在 SMEM 中进行零填充。
Merkle Tree 归约
CTA 中所有消费者完成块哈希后，链值存储在 smem_leaves 中，消费者协作归约这些链值，产出单个 32 字节根值。
完美 Merkle Tree：当叶子数是 2 的幂时（99% 的情况，除最后一个 CTA 外），归约是标准的迭代二叉树。每一层将活跃线程数减半，在读写之间同步以防止竞争。内部节点压缩使用 COMPRESS_PARAMS_INNER_NODE 参数，计数器为 0，块长度为 64，标志为 KEYED_HASH | PARENT。
不完美 Merkle Tree：当叶子数不是 2 的幂时，按照规范规定的结构处理。通过查看叶子数的二进制分解，对每个 2 的幂计算完美树，然后从右到左"串联"它们。例如 13 = 8+4+1 个叶子，得到相应的树结构。在 kernel 中，通过位扫描循环将线程分配到子树，每组线程独立地使用与完美树相同的减半算法归约其子树。最后，线程 0 从右到左依次归约所有子树根值。
Kernel 2 和 3：层级聚合
第一个 kernel 之后，每个 CTA 产出一个 32 字节根值。如果 CTA 很多，需要进一步归约。
ComputeBlakeMTKernel（Kernel 2）：从 GMEM 加载 kLeavesPerMTBlock 个根值（256、512 或 1024）到 SMEM，使用相同的 compute_perfect_mt / compute_blake_mt 例程归约。当网格只有一个块时，设置 IsSingleBlock 模板参数，在最终压缩上启用 ROOT 标志，将结果标记为最终哈希。
ReduceRootsKernel（Kernel 3）：如果 Kernel 2 产出超过一个块，Kernel 3 执行最终的单 CTA 归约；此 kernel 始终设置 ROOT 标志，因为之后不再有 kernel。
性能测试
我们将实现与一个流行的 BLAKE3 实现 BLAZE3 进行了对比测试。在最大的问题规模上，我们的实现比 BLAZE3 快约 370 倍。
未来工作
首先，由于 Pearl 挖矿不需要，我们尚未实现 BLAKE3 的 eXtendable-Output Function（XOF）模式，该模式允许生成任意长度的摘要。基于我们实现的未来版本可以添加此功能。
其次，我们提出 Merkle Tree 归约阶段有可能通过 warp 内部原语进一步优化。这些原语允许线程之间直接交换数据而不经过 SMEM，在树归约中很常见。此类原语也可用于让多个线程哈希同一个块。
结论
在本文中，我们展示了高性能 GPU 工作负载中常用的技术——如 TMA、swizzle 等——不仅适用于 GEMM、注意力等典型场景。我们将这些相同的模式应用于一个非常不同的问题：在单个大输入上高效计算密码学哈希。
我们的实现还表明理解被优化算法的重要性：例如，如果不观察到各块压缩是独立的，我们就无法利用其令人尴尬地并行的特性；如果不注意到可以部分融合 Merkle 归约，就必须执行不必要的 GMEM 往返。
所有这些设计选择加在一起，使我们在大问题规模上获得了接近 400 倍的显著加速。
一句话总结
Pearl 团队将 BLAKE3 哈希算法重写为单输入并行 GPU 架构，通过 TMA 加载、共享内存 swizzle、双流水线和 Merkle Tree 归约融合等技术，在大尺寸输入上实现了约 370-400 倍的速度提升。