面向AI ASIC上全同态加密NTT加速的低成本多精度脉动阵列

大家读完觉得有帮助记得关注和点赞摘要—全同态加密FHE可提供极强的数据隐私保障但存在极高的计算开销。利用张量处理器TPU等AI硬件加速FHE前景广阔却受限于根本性的精度不匹配问题TPU针对8位算术优化而FHE及其核心组件如数论变换NTT需要高精度运算。现有方案通过矩阵分解在低精度矩阵引擎上执行NTT计算但全精度结果的重构需要移位累加操作与矩阵乘法的数据流不匹配被迫将全精度重构从矩阵引擎卸载到向量处理器打断了矩阵乘法的数据流形成严重的性能瓶颈。为解决这一局限我们提出一种轻量修改的多精度脉动阵列可在阵列内部以统一数据流同步完成低精度矩阵乘法与原生全精度输出重构。采用OpenROAD在7nm工艺下综合的结果显示该设计引入的硬件开销可忽略不计。基于SCALE-Sim的周期精确仿真表明在128×128矩阵引擎上所提架构原生执行NTT可实现至少1.33倍的加速变换尺寸覆盖2¹²至2¹⁶成功使标准AI硬件支持高精度FHE加速。I. 引言全同态加密的核心安全优势在于支持直接在密文上执行任意计算全程无需暴露明文。尽管优势显著但其极高的计算开销传统上依赖定制化专用硬件[1][2]。然而这类定制加速器的设计与量产成本极高且难以规模化部署。更具可行性的替代方案是利用现有AI加速器如TPU——这类硬件已具备大规模并行能力与高能效特性[6]复用现有平台是缩小FHE落地性能差距的务实路径[2]。将FHE工作负载映射到AI加速器时存在根本性的架构不匹配TPU针对8位整型算术优化[6]而FHE需要基于大整数模的算术运算[7]。尽管FHE可通过残基数系统RNS[8]将大整数拆分为30~60位的独立通道但这些分量仍超出TPU矩阵引擎的原生支持范围。现有通用解决方案是将高精度数值拆分为多个低精度如8位矩阵运算再将运算结果聚合为全精度值。该聚合过程需要一系列移位累加操作与数位级进位传播而现有TPU架构无法在脉动矩阵引擎内完成这些操作被迫将计算卸载到外部向量单元。这种“外部累加”会引入大量数据搬运开销打断矩阵乘法的数据流导致矩阵引擎利用率低下最终抵消加速器的性能优势。本文证明仅需极小的硬件改动即可在AI加速器上高效执行FHE。我们对脉动阵列最后一行的处理单元PE做轻量修改使其支持矩阵引擎内部的原位移位加操作。不同于现有依赖可重构SIMD风格脉动阵列的多精度矩阵计算方案[9]我们的设计引入的硬件开销可忽略同时原生支持多精度累加。通过将重构操作嵌入矩阵乘法的统一数据流我们消除了外部累加的需求实现了加速器利用率的持续高位运行。该方法弥合了FHE的高精度需求与现代AI硬件的低精度效率之间的鸿沟可在兼容现有工作负载的前提下实现FHE的实用化部署。本文的主要贡献如下我们识别出将FHE核心操作如用于加速多项式乘法的NTT映射到AI加速器时的关键性能瓶颈全精度重构阶段尤其是移位累加与数位级进位传播。我们提出一种轻量修改的脉动阵列可在统一数据流下同步完成低精度矩阵乘法与内部全精度重构。7nm工艺OpenROAD实现结果显示该双操作支持的硬件成本增量低于1%。我们采用SCALE-Sim周期精确模拟器评估性能对于占FHE运行时主导地位的NTT在变换尺寸2¹²至2¹⁶范围内所提架构实现至少1.33倍的加速同时兼容传统低精度AI矩阵引擎引入的开销极小。II. 背景与相关研究FHE方案的核心运算是线性代数与多项式算术[11]。密文表示为环元素的向量密钥切换、自同构、密文乘法等核心操作最终都可归约为结构化线性变换与多项式卷积[3][4][11][5]。近期研究证明这类内核可被重构为密集矩阵运算从而在GPU与AI加速器上实现加速。TensorFHE[3]将FHE原语适配到GPU张量核TensorFHE[4]进一步通过线性代数抽象提升了跨异构加速器的可移植性。CROSS[5]通过将高精度模运算重构为低精度矩阵乘法验证了将FHE工作负载映射到脉动阵列加速器的可行性该转换对NTT这类FHE核心操作[12]尤为关键——NTT可大幅简化多项式乘法。A. 矩阵化NTT及其到AI ASIC的映射定制FHE加速器通常依赖专用并行蝶形单元执行NTT[2]。而面向AI ASIC的方案则复用现有矩阵乘法引擎与向量处理器FHE应用所需的大尺寸NTT2¹²至2¹⁶个元素通过4步NTT算法[13]分解为一系列矩阵友好的子操作。对于大小为Nr×c的一维输入4步NTT的执行流程如下列变换沿矩阵列方向执行c次独立的r点NTT。该操作等效为预计算的r×r变换矩阵与r×c数据矩阵的矩阵乘法。旋转因子乘法为修正二维转换引入的相位偏移中间结果与预计算的r×c旋转因子矩阵逐元素相乘。矩阵转置将r×c矩阵转置为c×r布局将原行数据对齐为连续的列。行变换沿转置后的新列方向执行r次独立的c点NTT。该操作等效为第二次矩阵乘法使用c×c变换矩阵。最终4步分解将一个N点NTT替换为两次N​×N​矩阵乘法、一次逐元素乘法与一次转置操作。这种形式与现代AI加速器中的脉动矩阵单元高度契合脉动阵列可原生执行矩阵乘法辅助处理单元负责转置与逐元素操作。B. FHE与AI ASIC的精度不匹配将FHE工作负载映射到AI ASIC时存在显著的架构不匹配FHE需要基于大整数模的算术通常跨度数百比特而现代AI ASIC依赖8位整型单元以最大化能效与可扩展性典型配置为128×128及以上规模的PE阵列[6]。为弥合精度差距现有方案如CROSS的优化映射[5]将高精度矩阵运算拆分为多个低精度子步骤再聚合中间结果恢复目标精度。但该重构/聚合阶段成为关键性能瓶颈矩阵乘法完成后执行流程被迫串行化且重构操作需卸载到主矩阵引擎外的向量处理器引入数据搬运开销降低脉动阵列利用率。这种低效性正是本工作的出发点我们将高精度重组直接嵌入矩阵乘法的数据流增强标准8位脉动阵列以支持重构步骤无需阵列外计算即可得到全精度矩阵乘法结果。值得注意的是我们的方案不依赖传统的可重构SIMD多精度架构[9][14]——这类架构通过分组相邻PE模拟更宽的算术单元虽在其他领域有效但会引入显著的硬件开销且随着精度要求提升可用于并行矩阵乘法的有效阵列规模会大幅缩减。III. 增强型脉动阵列中的多精度矩阵乘法TPU矩阵引擎脉动阵列由小位宽乘法器与较宽累加器构成。例如谷歌TPU的每个PE内置8位整型乘法器而权值静止数据流下的列向累加归约采用32位加法器[6]这种设计使其在低精度密集线性代数运算中效率极高。但高精度运算输入操作数超过8位无法在阵列内直接执行大整数必须拆分为更小的数位digit在对应的低精度元素上完成矩阵乘法后再通过移位加操作聚合中间结果得到最终的全精度结果。现有脉动阵列不支持内部原生全精度重构该步骤通常卸载到阵列外的向量处理器打断了原本统一的矩阵乘法数据流引入不必要的数据搬运最终降低矩阵引擎的利用率。为解决这一问题我们首先说明如何将全精度输入矩阵分解为仅由小digit如8位整型构成的矩阵再描述标准矩阵乘法如何生成中间结果最后说明如何在脉动阵列内处理这些中间结果将其作为原始数据流的无缝扩展恢复正确的全精度输出。图1通过分解与重构实现高精度矩阵乘法到低精度矩阵乘法的转换A. 高精度矩阵乘法流程概述图1展示了含多digit高精度整数的矩阵乘法如何转换为等价的单digit操作数矩阵乘法可在脉动阵列上高效执行。示例采用两位十进制数以简化说明该流程适用于任意进制。图顶为原始高精度矩阵乘法的计算结果。首先每个全精度矩阵被分解为仅由更小digit构成的更大矩阵图1步骤1左矩阵的每个元素被展开为一个行向量从最低有效位到最高有效位依次放置到连续列中得到更宽的单digit矩阵右矩阵的转换方式不同其组成digit被排列为类托普利兹结构[15]编码笔算digit乘法所需的所有digit交叉乘积[16]——右矩阵的每个digit被复制到多列并做适当的水平移位确保单digit矩阵相乘时digit对的乘积自动对齐到正确的位置权重。最终单次密集矩阵乘法即可生成原始高精度计算所需的所有digit交叉乘积。以左矩阵元素a、右矩阵元素b为例说明两者均为两位十进制数可表示为aa1​⋅10a0​bb1​⋅10b0​其中a1​,a0​,b1​,b0​为单个digit。乘积ca×b可展开为二次多项式该计算可表示为矩阵向量乘法如图1所示分解后的digit采用最低有效位优先的排列顺序这对脉动阵列内的操作调度至关重要详见下一节。单digit矩阵相乘得到2×6的中间结果矩阵图1步骤2随后需要通过重组合并中间结果得到正确的全精度输出。如图1步骤3所示重组过程需要将对应不同digit位置的相邻列通过移位与数位级进位传播合并完成多digit加法。需注意进位仅局限在每个digit组内不会扩散到其他输出digit对应的列。B. 所提脉动阵列架构要在脉动阵列内原生完成精度重构且不打断典型数据流几乎无需修改标准操作仅需为脉动阵列最后一行的PE增加一个额外的加法器其余阵列结构保持不变。我们的架构基于标准谷歌TPU风格脉动阵列采用广泛使用的权值静止数据流右矩阵权值预加载到阵列中左矩阵从西侧边缘按行水平流入阵列。每个PE执行流入数据与本地存储权值的8位乘法16位乘积通过32位加法器沿每列垂直向下连续累加典型结构如图2所示。最后一行的PE被称为重构处理单元RPE相比其余行的普通PE新增一个输入端口接收左侧RPE传来的digit并将其加到同列向下传播的求和结果上。这种设计使数位进位沿阵列最后一行水平传播与每列求和结果的垂直传播完全同步。当结果超出全精度输出对应的列范围时多路选择器将为加法器选择零输入停止进位传播。图2所提脉动阵列架构。所有PE保持典型的乘加结构仅最后一行PERPE新增一个加法器将水平数位级进位传播与垂直列向累加结合。图3展示了矩阵乘法的脉动波前传播过程中与列向垂直归约同步执行的重构操作的按周期运行示例。假设图1分解得到的右矩阵已预加载到阵列中左矩阵从左向右逐行流入。为聚焦重构操作图3仅考虑图1中分解后左矩阵的第一行。随着脉动阵列运行中间结果矩阵的第一行逐渐在阵列底部生成从第3个周期开始连续列依次输出首个结果这些结果进入最后一行的RPE。每个RPE保留本列的最低有效字节将最高有效字节作为进位传给右侧相邻的RPE。这种数位级进位传播持续到全精度结果完全重构。需注意中间结果以最低有效位优先的顺序流入该顺序至关重要确保第k位digit在第k−1位digit在同列产生进位之后才到达。图3单digit矩阵乘法所需的列向垂直归约与全精度矩阵乘法输出的水平数位级进位重构同步执行的运行示例。C. NTT专用的全精度矩阵分解本节介绍针对NTT的全精度矩阵分解方法结合对齐基变换BAT[5]进行简化。4步NTT将N点变换分解为更小的子变换计算为N×N输入矩阵A与旋转因子矩阵T的乘积。假设A和T包含b位全精度整数要在仅支持w位乘法器wb的所提脉动阵列上执行该乘法需采用III-A节的方法分解两个矩阵每个元素拆分为k⌈b/w⌉个radix-2w的digit对应位置权重20,2w,…,2(2k−2)w。该扩展会将矩阵A的尺寸从N×N调整为N×kN将矩阵T的尺寸从N×N调整为kN×(2k−1)N。矩阵T仅包含依赖变换参数(N,q)的旋转因子对特定NTT实例而言是运行时常数。BAT可利用这一特性缩减T的尺寸分解后的T的每个digit的贡献需按模q加权因此所有超过2kw的digit权重可折叠回低位digit。由于T在设计阶段已知且运行时恒定这种digit的重分配可离线完成。以下玩具示例说明旋转因子digit的离线缩减过程将图1右矩阵的第一个元素54作为旋转因子之一采用两位精度十进制digit分解扩展为托普利兹矩阵假设所有计算模97权重102的digit 5的贡献可离线化简为5×102mod97151×1015×100。将该结果分配到权重100与101的列得到重复上述旋转因子digit折叠操作直到T的所有digit都适配w位为止。采用BAT建议的离线重分配后分解后的T的维度缩减为kN×kN。IV. 实验评估为评估所提方法我们量化了两个核心指标①采用4步NTT算法执行FHE NTT内核变换尺寸2¹²至2¹⁶的性能加速比②相比标准TPU的硬件开销增量。性能测试采用周期精确的SCALE-Sim模拟器[17]建模架构与商用AI加速器类似模拟基线包含128×128脉动矩阵单元MXU与128车道向量处理单元VPU[6]。每个MXU集成8MB片上输入、权值与输出缓冲区带宽足以支撑峰值脉动吞吐量每个周期为每个PE提供一个操作数如128×128阵列每周期交付128个8位字确保阵列持续满载。我们对比了两种现有方案①TensorFHE[3]在MXU上执行低精度矩阵乘法全精度重构卸载到VPU②CROSS[5]矩阵乘法与结果重构方法与TensorFHE一致但原始矩阵到低精度矩阵的分解方式不同。所有方案包括所提架构在矩阵乘法尺寸超过MXU规模时均采用分块执行策略[18]。由于FHE方案普遍采用32~64位模运算对应格密码学与主流FHE库最常用的模数我们评估了这两种配置的性能。A. 运行时性能对比表I总结了32位与64位全精度输入下TensorFHE、CROSS与所提架构执行完整4步NTT的周期数对比。所有架构的执行时间均分为MXU周期对应8位digit矩阵乘法与VPU周期对应脉动阵列外的所有操作TensorFHE与CROSS的VPU周期包含重构全精度结果所需的移位加操作、模规约、旋转因子乘法以及4步NTT涉及的矩阵转置。所提架构的VPU周期仅包含模规约因为全精度重构由阵列内的RPE在矩阵乘法过程中在线完成。表I32位/64位全精度输入下不同变换尺寸时TensorFHE[3]、CROSS[5]与所提架构的4步NTT算法运行时对比周期数及相对加速比NTT尺寸TensorFHE[3] MXU周期TensorFHE[3] VPU周期TensorFHE[3] 总周期CROSS[5] MXU周期CROSS[5] VPU周期CROSS[5] 总周期所提架构 MXU周期所提架构 VPU周期所提架构 总周期21214240217616416356621765742356625638222131526442881955291744288134629174448962221416288857624864163188576248941631889617214215428161702459840428461702459328428461664445102168163234048115680816623404811571081662332884990(a) 32位数据宽度NTT尺寸TensorFHE[3] MXU周期TensorFHE[3] VPU周期TensorFHE[3] 总周期CROSS[5] MXU周期CROSS[5] VPU周期CROSS[5] 总周期所提架构 MXU周期所提架构 VPU周期所提架构 总周期2125696083206528014270832022590142702561452621361056165767763236702165765327836702448371502146515233152983046527833152984306527889666174215171264661762374401713906617623756617139016641730542163265281323524588803266541323524590063266543328329982(b) 64位数据宽度实验结果显示两种精度设置下所提架构均能减少执行周期加速比范围为1.33×~4.49×32位与64位配置的性能趋势与加速比几乎一致。性能提升的核心来源是消除了基于VPU的全精度输出重构TensorFHE[3]与CROSS[5]的非可忽略比例的运行时被用于阵列外的部分和累加与数位级进位传播。而我们的架构将这些操作直接嵌入MXU数据流完全消除了这部分开销。最小NTT尺寸下的加速比最高例如32位配置下212点NTT的执行时间从16416周期降至3822周期提升达4.30×。这是因为小尺寸下VPU执行占总运行时的比重较高消除后收益显著。随着NTT尺寸增大矩阵乘法逐渐成为运行时主导消除VPU开销的相对收益下降最大216点NTT的加速比收敛到1.36×。B. 物理综合对比为评估RPE新增的加法器与多路逻辑的开销我们采用OpenROAD物理综合平台与ASAP7工艺库对基线设计与所提设计进行实现目标时钟频率为1GHz匹配工业级TPU规格[6]。图4展示了小规模8×8所提脉动阵列的物理布局示例。图4所提8×8脉动阵列的物理布局。表II报告了不同阵列规模下的后布局面积与功耗以及新增全精度输出重构逻辑带来的相对开销。结果显示所提修改引入的硬件代价极低面积开销从8×8阵列的2.79%下降到256×256阵列的0.09%功耗开销从8×8阵列的5.52%下降到256×256阵列的0.17%。这种缩放趋势源于随着阵列总规模增大最后一行新增逻辑的开销被充分摊销。表II采用OpenROAD物理综合平台在7nm工艺下实现的传统脉动阵列与所提增强型最后一行含RPE设计的面积与功耗对比MXU尺寸基线面积(mm²)所提设计面积(mm²)面积开销基线功耗(W)所提设计功耗(W)功耗开销8×80.00410.00422.79%0.02540.02415.52%16×160.01650.01671.39%0.09890.09632.76%32×320.06590.06640.70%0.39030.38501.38%64×640.2630.2640.35%1.841.860.69%128×1281.0541.0560.17%7.397.410.34%256×2564.2174.2210.09%29.5629.610.17%总体而言所提方法通过大幅降低延迟仅引入可忽略的功耗开销显著减少了单次NTT的总能耗证明脉动阵列内的原生全精度输出重构既是有效的性能优化也是实用的高能效硬件增强方案。V. 结论在AI ASIC上执行FHE内核如NTT是比定制硬件更易获取、能效更高的替代方案但核心矛盾在于AI ASIC针对低精度算术优化而FHE需要高精度运算。现有方法采用标准数据流将高精度工作负载映射到低精度矩阵引擎时会在全精度结果重构阶段产生串行瓶颈。我们提出的新型多精度脉动阵列架构将低精度矩阵乘法与原位高精度重构统一起来设计了全新的数据流矩阵乘法结果垂直流动与水平的数位级进位传播完美同步。这种融合方法消除了重构瓶颈为NTT内核带来了显著的执行加速硬件开销不足1%。