1. 项目概述与核心挑战零知识证明ZKP正在从理论走向大规模应用成为可验证计算和隐私保护区块链的核心基石。然而其证明生成过程尤其是核心的椭圆曲线多标量乘法MSM和数论变换NTT计算开销巨大严重制约了其落地。例如为一个ImageNet ViT模型生成证明可能需要近一小时。传统上加速这些运算依赖于GPU、FPGA甚至专用ZKP ASIC。但近年来AI专用集成电路ASIC如谷歌的TPU因其在矩阵运算上提供的极致计算密度和能效比展现出了新的潜力。问题在于TPU这类AI加速器的硬件架构如巨大的矩阵乘法单元MXU、粗粒度的向量寄存器VReg与ZKP算法如精细的位操作、不规则的数据访问模式之间存在根本性的不匹配。直接移植现有算法会导致性能急剧下降有时甚至慢30倍以上。MORPH框架的诞生正是为了系统性地解决这一“硬件-算法”鸿沟。它不是一个简单的端口而是一次从底层算术到高层数据流的全面重构。其核心思想是与其让算法去适应硬件的缺陷不如重新设计算法使其计算模式完美契合AI ASIC特别是TPU的硬件特性。这涉及到两个层面的深度优化在算术层面将高精度素数域上的模运算转化为TPU MXU最擅长的低精度、密集的矩阵乘法GEMM在数据流层面重新设计MSM和NTT的计算流程彻底避免TPU上极其昂贵的细粒度数据重排和跨芯片通信。最终目标是将TPU这类原本为AI训练设计的“蛮力”计算引擎改造为高效的ZKP证明生成加速器。2. 硬件瓶颈分析为什么传统ZKP算法在TPU上“水土不服”要理解MORPH的优化必须先看清TPU与ZKP算法之间的根本矛盾。TPU是一种高度特化、追求极致吞吐量和能效的硬件。如图2所示其计算核心围绕向量寄存器VReg展开每个VReg是8x128的32位值块所有计算单元VPU向量处理单元、MXU矩阵乘法单元、XLU布局转换单元都以VReg为粒度进行读写。这种设计带来了高效率但也牺牲了灵活性。2.1 TPU编程模型的约束计算单元异构性MXU矩阵乘法单元核心算力来源一个巨大的脉动阵列如128x128专为bf16/int8矩阵乘法优化。要发挥其峰值性能计算必须组织成对齐的矩阵块如8x128x128。小矩阵或不规则形状会导致严重的利用率不足。VPU向量处理单元执行32位标量/向量运算并行度PAR_VPU可达2048 ops/cycle。同样向量长度必须是VReg粒度1024个元素的倍数才能满载。XLU布局转换单元负责VReg之间的数据重排转置、广播、洗牌。对于VReg粒度的操作效率很高但对于元素级的随机访问即“洗牌”成本可能高达每个元素一个周期这在ZKP的蝴蝶网络等操作中是灾难性的。内存层次与带宽片外高带宽内存HBM带宽比片上VReg间带宽低1-2个数量级。因此高性能TPU内核必须最大化片上数据复用并最小化与HBM的数据交换。2.2 传统ZKP算法的“阿喀琉斯之踵”当我们将为GPU/CPU设计的SOTA ZKP算法直接移植到TPU时会触发上述所有硬件约束算术层面ZKP运算基于大素数域如256、377、753位。硬件原生支持32位因此需要将一个大整数分解为多个32位“数字”进行运算。传统的基数Radix蒙哥马利约减算法需要进行O(D^2)次32位乘法D是数字个数并伴随O(D)步的顺序进位链传播。在TPU上进位链要求频繁的、细粒度的数据依赖和重排这正好撞上了XLU的短板导致计算单元MXU/VPU大量空闲等待数据重组。数据流层面MSMPippenger算法其“预排序”阶段涉及根据标量值对海量椭圆曲线点进行随机散射scatter和聚集gather。在GPU上这由成千上万个CUDA核心并行处理。在TPU上这种细粒度的、不规则的内存访问模式会导致大量的XLU开销和低效的HBM访问成为主要瓶颈。NTT蝴蝶算法经典的Cooley-Tukey蝴蝶算法计算复杂度最低O(N log N)但每一级都需要对远小于VReg粒度的数据块进行洗牌。这相当于要求TPU的XLU进行海量的、低效的元素级置换完全无法发挥其计算优势。关键洞察传统的算法复杂度分析如大O表示法只关心操作总数隐含假设了一个同构的、细粒度可控的硬件如CPU。对于TPU这种异构的、粗粒度的并行硬件它完全忽略了不同计算单元间的瓶颈差异和布局转换的巨额开销。因此一个理论上操作数更少的算法在TPU上的实际运行时可能远慢于一个操作数更多但计算模式更“规整”的算法。3. 方法论基石Big-T复杂度模型为了量化并指导在AI ASIC上的算法设计MORPH提出了Big-T复杂度模型。这是一个硬件感知的性能下界模型。T(N) O( max( max_k(W_k / P_k), Mem ) )其中W_k(N)映射到第k类计算单元如MXU, VPU, XLU上的总工作量。P_k第k类计算单元的峰值并行度如图2中的PAR_MXU,PAR_VPU,PAR_S。max_k(W_k / P_k)捕获了异构流水线计算单元中的瓶颈阶段。例如即使MXU很快如果VPU或XLU是瓶颈整体速度仍受限于最慢的阶段。Mem捕获片外内存访问的总体延迟。Big-T模型的意义在于它迫使算法设计者不仅要减少总操作数更要平衡各硬件单元的工作负载并最小化昂贵的操作如XLU洗牌和HBM访问。MORPH后续的所有优化都以此模型为指导旨在降低Big-T复杂度中的主导项。4. 算术优化MXU-centric RNS惰性约减这是MORPH解决高精度模运算与低精度硬件之间矛盾的核心创新。4.1 问题根源素数域与进位链ZKP在素数域F_M上进行运算模数M是一个大素数如2^256。素数无法被分解为互质因子因此无法直接使用余数系统RNS的优势。RNS可以将一个大数用一组小模数的余数表示在该系统内加法和乘法可以独立在每个小模数上并行进行无需进位。但最终如果需要还原回原始整数或进行模M约减则需要复杂的“基转换”成本很高。因此SOTA的GPU方案采用基数蒙哥马利约减。它将一个M以内的数表示为D个32位数字基数2^32。乘法a * b mod M需要大整数乘法计算a*b结果最多2D个数字需要O(D^2)次32位乘法和一次O(D)的进位传播。蒙哥马利约减将2D位的中间结果约减回D位又需要O(D^2)次乘法和一次O(D)的进位传播。这两次O(D)的进位链是顺序的、数据依赖的在TPU上会转化为大量的XLU操作成为Big-T模型中的主要瓶颈见表1。4.2 解决方案扩展域与MXU加速的惰性约减MORPH的算术优化分为两步精妙的转换第一步逃逸到扩展域既然在素数域F_M内无法避免进位MORPH选择将整个计算暂时“搬”到一个更大的非素数域F_Q中。我们精心选择一个复合数Q使其满足Q M^2并且Q可以分解为一组小的、互质的模数{q_i}。这样F_Q就拥有了一个自然的RNS表示。操作将输入x, y ∈ F_M通过中国剩余定理CRT映射到它们在F_Q中的RNS表示x_Q, y_Q。每个x_Q现在由一组比如2D个小余数表示。优势在F_Q的RNS中乘法x_Q * y_Q变得极其简单只需对每一对对应的小余数进行独立的模q_i乘法共O(2D)次操作完全消除了进位链。因为Q M^2所以两个域内数的乘积在F_Q中也不会溢出我们暂时不需要对Q进行模约减。第二步MXU-centric RNS惰性约减现在我们有z_Q x_Q * y_Q在F_Q中。我们需要将结果映射回F_M即计算z (x * y) mod M。传统的做法是RNS重建得到大整数z_Q→ 模M约减 → 再次RNS分解。这个过程成本很高。MORPH的关键创新在于它利用Simon的约减算法[24]和基对齐变换Basis Aligned Transformation[35]将上述多步过程压缩为一次低精度矩阵乘法和几次向量操作。如算法1和图3所示字节分解将x_Q的每个32位余数称为“肢”limb分解为4个字节。这使得数据可以被组织成uint8类型。核心矩阵乘法执行一个uint8类型的矩阵乘法算法1第19行Einsum。这个操作完美匹配TPU MXU的强项。矩阵[E]是预计算的编码了从F_Q到F_M约减所需的常数。向量操作伴随矩阵乘法的是一些32位的向量点积、移位和加法算法1第16, 17, 20, 21行这些由VPU高效执行。最终效果原本O(D^2)的蒙哥马利约减计算其主导部分被转化为一个可以被MXU高度并行化、吞吐量极高的uint8矩阵乘法。虽然内存占用因RNS表示而增加了约2倍但计算瓶颈从顺序的、洗牌密集的XLU操作转移到了并行的、计算密集的MXU/VPU操作。如表1所示Big-T复杂度中的XLU跨度从O(D^2 log D)降到了近乎0而MXU跨度虽然仍是O(D^2)但因其极高的PAR_MXU而被有效分摊。实测中这带来了高达90倍的算术加速。实操心得参数选择与权衡选择扩展域模数Q是关键。Q必须大于M^2以确保中间结果不溢出同时其因子分解{q_i}应尽可能小且数量适中以平衡RNS的并行度和计算/存储开销。通常q_i会选择接近机器字大小如32位的素数使得模q_i的运算在硬件上高效。这需要针对具体的TPU型号如v4, v5p, v6e和精度256/377/753位进行微调。5. 数据流优化一布局固定的Pippenger MSMMSM计算Σ (scalar_n * point_n)。Pippenger算法通过将标量按比特窗口分组预排序相同窗口值的点进行累加来减少点加次数。但将其移植到TPU面临三大挑战1) 跨TPU通信2) TPU内布局重组3) 从HBM重复加载点数据。5.1 LS-PPG算法详解MORPH提出了布局固定的PippengerLS-PPG其核心思想是在预排序和桶累加这两个主要阶段强制执行统一的数据分片和内存布局使得前一阶段的输出可以直接作为下一阶段的输入无需重排。算法2和图4展示了LS-PPG的流程统一分片将MSM的N个点对(scalar, point)沿窗口维度K进行分片每个TPU核心负责处理一个窗口子集的所有点。这样每个核心内部点的预排序是独立的。桶化与预排序融合在每个TPU核心上算法直接根据标量的当前窗口值将点“散射”到对应的桶bucket中。这个过程在内存中产生的是已经按照桶索引组织好的数据布局。布局固定的桶累加由于数据已经在桶中有序排列桶累加阶段可以直接对每个桶内的连续点序列进行向量化点加运算完全避免了在预排序后需要进行的全局数据重排reshape/shuffle。树形桶归约传统的桶归约是顺序的。LS-PPG将其改为树形归约将2^c个桶的归约深度从O(2^c)降低到O(c)暴露了更多并行性提高了VPU利用率。窗口合并最后按窗口顺序进行标量乘和点加得到最终结果。5.2 性能收益分析LS-PPG带来的改进直接体现在Big-T模型中表2XLU跨度与基线预排序PPG相同但基线算法在TPU上实现时其O(2^c * N log N)的洗牌成本会被PAR_SXLU的洗牌并行度严重惩罚。LS-PPG通过避免额外的布局转换在实际中大幅减少了XLU活动。内存跨度这是最大的改进。基线算法需要为每个窗口从HBM重复加载所有N个点导致内存访问量为O(K * N)。LS-PPG通过融合预排序和桶累加使得每个点在整个MSM计算中只被加载一次内存访问量降至O(2N)。对于大问题规模这直接转化为数倍的性能提升。注意事项桶大小与负载均衡LS-PPG的桶化阶段可能产生负载不均衡因为不同桶内点的数量可能差异很大。在实践中需要设置一个每桶最大点数N并进行动态调度或填充。过小的N会限制并行度过大的N则浪费内存。需要根据具体的椭圆曲线点和TPU内存容量进行权衡。6. 数据流优化二五步NTTNTT是多项式承诺方案中的核心计算X_k Σ (x_n * ω^{kn}) mod M。蝴蝶NTT在TPU上因细粒度洗牌而失败。布局不变的3步NTT[35]通过将一维NTT重构为二维矩阵运算消除了洗牌但其性能受限于大型矩阵乘法GEMM的规模。6.1 从三步到五步递归分解以降低参数存储3步NTT将长度为N R * C的向量重塑为R x C矩阵然后进行对每一行做长度为R的NTT需要R x R的旋转因子矩阵TF_R。逐元素乘以R x C的旋转因子矩阵TF_C。对每一列做长度为C的NTT需要C x C的旋转因子矩阵TF_C。问题在于步骤1所需的TF_R矩阵大小为R^2。当N很大时如2^24R约为2^12TF_R矩阵将占用巨大内存2^24个元素可能导致内存溢出OOM。MORPH的5步NTT图5c通过递归分解来解决这个问题。它将R进一步分解为R R1 * R2从而将步骤1中的行NTT一个R点的NTT本身再用一个3步NTT来实现步骤1a将R1 x R2的行数据视为R1 x R2矩阵对行做R1点NTT。步骤1b逐元素乘旋转因子。步骤1c对列做R2点NTT。这样原本需要存储的巨型R x R矩阵TF_R被分解为三个小得多的矩阵TF_{R1}(R1 x R1),TF_{R2}(R2 x R2), 和一个中间的旋转因子矩阵。总存储开销从O(R^2 C^2)降至O(R1^2 R2^2 C^2)。6.2 计算与通信的权衡5步NTT增加了计算步骤从3步到5步但如公式1所示它通过将一次大的GEMMR x R分解为多次较小的GEMMR1 x R1,R2 x R2可能更好地适应TPU MXU的切片大小提高计算效率。更重要的是它解决了大规模NTT的参数存储问题。在Big-T模型表2中MXU跨度3步NTT为N*(RC)/PAR_MXU5步NTT为N*(R1R2C)/PAR_MXU。通过选择R1 ≈ R2 ≈ sqrt(R)可以降低MXU的工作负载。内存跨度5步NTT略高因为多了两次数据重塑reshape操作但仍在可接受范围。XLU跨度两者都避免了蝴蝶NTT的O(N log N)洗牌只有固定的数据重塑开销。实验表明在较小规模N 2^20时3步NTT可能更优因为因子R1, R2太小无法充分利用VReg。当规模更大N 2^22时5步NTT凭借其可管理的内存占用和良好的MXU利用率开始显现优势。7. 实现、评估与实战经验MORPH使用JAX实现利用其XLA编译器为TPU生成高效代码。评估在TPUv6e-88个芯片上进行并与在NVIDIA V100上运行的SOTA GPU库GZKP进行对比。7.1 性能结果NTT加速显著对于753位高精度NTTTPUv6e-8相比V100实现了最高10倍的吞吐量提升。这主要归功于MXU-centric RNS算法将高精度运算高效映射到了MXU上以及数据流优化彻底消除了洗牌瓶颈。TPU的性能随精度增加而下降的幅度1.3-3倍远小于GPU6-7倍显示了其对高精度运算的良好扩展性。MSM竞争力对于753位MSMTPUv6e-8达到了与GZKP (V100)可比的吞吐量1.2倍。瓶颈主要在于桶化阶段的随机散射操作这对内存带宽不友好。对于256位MSMTPU性能低于GPU说明细粒度的数据移动仍是TPU的弱点。批处理优势如图7所示MXU-centric RNS惰性约减在批处理场景下优势巨大。随着批处理规模增大MXU的利用率得到充分发挥加速比可达157倍。NTT也需要一定的批处理128来填满VReg达到峰值性能。7.2 常见问题与调优指南在实际部署MORPH或类似优化时可能会遇到以下问题精度与溢出在扩展域F_Q中运算必须确保Q M^2 * N对于N次连续乘法否则会因溢出导致错误。需要根据具体电路的最大乘法链长度来安全地选择Q。TPU内存限制RNS表示会使内存占用增加约2倍。对于超大规模问题如N 2^24即使使用5步NTT也可能面临OOM。此时需要结合模型并行将数据分片到多个TPU芯片上。JAX/XLA编译开销TPU上的JAX代码首次运行需要较长的编译时间。对于生产环境建议提前编译AOT好计算图或使用持久化编译缓存。选择3步还是5步NTT这没有绝对答案。一个实用的启发式方法是对于N 2^20优先尝试3步NTT对于更大的N如果3步NTT因参数矩阵过大导致OOM则切换到5步NTT。可以编写一个简单的性能探查器在小规模数据上测试两种方法的实际性能。MSM桶大小c值选择在LS-PPG中窗口大小c决定了桶的数量2^c。c值越大桶越多预排序的碰撞概率越高点加次数越少但桶归约的开销和内存占用也越大。通常需要在具体硬件上做参数扫描c取值在12到20之间是常见范围。8. 总结与展望MORPH的工作清晰地展示了一条路径通过硬件感知的算法重构可以将为矩阵运算而生的AI ASIC转化为高效的ZKP加速器。其贡献不仅是具体的MSM/NTT优化更在于提供了Big-T这样的分析工具和“面向MXU编程”的设计哲学。从更广阔的视角看这项工作预示着专用硬件设计范式的转变。未来的ZKP加速器或许不必从头设计而是可以借鉴TPU这类成熟、高能效的AI加速器架构通过类似MORPH的编译器和算法层创新使其适配密码学计算。同时这也对AI加速器的设计提出了新要求例如增加对细粒度、不规则数据访问如散射/聚集的硬件支持将能进一步释放其在更广泛算法上的潜力。对于开发者而言MORPH的开源代码提供了一个极佳的起点。在尝试将其他ZKP原语如多项式求值、KZG承诺移植到TPU上时应始终牢记两个核心原则将计算转化为规整的矩阵/向量运算以及最大限度地保持数据布局的稳定性避免昂贵的重排。这或许是解锁下一代高性能隐私计算和区块链扩展性的关键。