💻 computer science

KeyMemRT Compiler and Runtime: Unlocking Memory-Scalable FHE

KeyMemRT 是一个基于 MLIR 的编译器和运行时框架，它利用数据流分析来自动管理全同态加密（FHE）旋转密钥的生命周期，与现有的最先进编译器相比，显著降低了内存消耗并提高了执行速度。

原作者： Eymen Ünay, Björn Franke, Jackson Woodruff

发布于 2026-01-27

📖 1 分钟阅读☕ 轻松阅读

原作者： Eymen Ünay, Björn Franke, Jackson Woodruff

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正试图解决一个巨大且复杂的拼图，但由于戴着眼罩，你只能在看不见的情况下操作拼图碎片。你必须在遵循一套特殊规则的前提下操纵这些碎片，而这些规则会将拼图隐藏起来，不让外人窥视。这本质上就是**全同态加密（Fully Homomorphic Encryption, FHE）**所做的事情：它允许计算机在数据保持完全加密的状态下进行计算。

然而，这里有一个巨大的陷阱。为了进行这种“盲视”数学运算，计算机需要一个庞大的特殊“魔法钥匙”库（称为旋转密钥/rotation keys）。

问题所在：“钥匙囤积者”

把这些魔法钥匙想象成一个拥有数千间客房的酒店里那套庞大的实体钥匙收藏。

旧方法 (ANT-ACE)： 想象一位酒店经理，在客人甚至还没到达之前，就把大楼里的每一把钥匙都抓了出来，堆在服务台前。即使客人只需要进入 101 号房，经理也会把 102 号到 5000 号房的钥匙全部放在桌上。
- 结果： 桌面（内存）变得杂乱无章且溢出。如果这家酒店规模宏大，桌子会被这些钥匙塞得满满当当，以至于没有空间放置其他任何东西。系统会崩溃或变慢，因为它正试图管理一座由不必要钥匙组成的“大山”。
“缓慢”的方法 (Fhelipe)： 另一位经理为了节省空间，只保留了一些“万能钥匙”。为了打开 101 号房，他手头没有直接的钥匙，所以必须先用一把万能钥匙打开 1 号房，再用另一把打开 2 号房，以此类推，通过一连串的钥匙进行链式操作，直到到达 101 号房。
- 结果： 桌面很整洁，但客人必须等待很长时间，因为经理要在这一长串钥匙链中反复摸索。这个过程非常缓慢。

解决方案：KeyMemRT

该论文的作者构建了一个名为 KeyMemRT 的新系统。你可以把它想象成一个超级智能、自动化的礼宾服务，能够完美地管理这些钥匙。

它精准了解你的需求： 它不会一次性抓取所有钥匙，而是通过分析客人的行程（程序的代码），来查看他们究竟会访问哪些房间以及访问的顺序。
即时交付： 它保持桌面整洁。它只会在客人即将需要 101 号房钥匙的那一刻之前，才把特定的钥匙拿出来。
即时清理： 一旦客人离开 101 号房，礼宾人员会立即将那把钥匙收回并归位，从而为下一把钥匙腾出空间。
“预取 (Prefetch)”技巧： 为了确保客人永远不必等待，礼宾人员会在客人还在当前房间时，就在后台开始准备下一把钥匙。这在后台运行，因此整个过程感觉既快速又流畅。

他们取得了什么成就？

研究人员使用各种复杂任务（如识别医学扫描图像或金融数据）对这个新系统进行了对比测试。

内存节省： 与“钥匙囤积者”方法 (ANT-ACE) 相比，KeyMemRT 将所需的内存减少了 1.74 倍。这就像是把一整家酒店的钥匙装进了一个单肩包，而不是存放在一个仓库里。
速度： 与“缓慢链式”方法 (Fhelipe) 相比，KeyMemRT 快了 1.73 倍。它不仅节省了空间，而且工作效率更高，因为它不会在从零开始制作钥匙上浪费时间。

为什么这很重要？

目前，FHE 很难使用，因为它需要拥有海量内存（数百 GB）的计算机来存放这些密钥。这使得它对于许多实际应用场景来说既昂贵又不切实际。

KeyMemRT 扮演了翻译官和交通管制员的角色。它将复杂的、混乱的代码进行优化，使其可以在标准计算机上运行，而不需要仅仅为了存放钥匙就动用超级计算机。它让隐私保护计算（例如在不接触原始数字的情况下分析敏感的医疗或金融数据）变得更加可扩展且高效。

简而言之： 他们构建了一个智能系统，阻止了 FHE 程序囤积无用的钥匙，也阻止了它们在从零开始制作钥匙上浪费时间，从而让隐私计算能够在更小、更便宜的机器上运行得更快。

技术摘要：KeyMemRT 编译器与运行时

问题陈述

全同态加密（FHE）能够实现对加密数据的隐私保护计算，但由于极高的计算和内存需求，面临着难以实现大规模实际应用的障碍。一个主要的瓶颈是**旋转密钥（rotation keys）**的管理。在分批次（batched）FHE 方案（如 CKKS、BGV、BFV）中，旋转操作对于向量内计算和自举（bootstrapping）至关重要。然而，这些方案要求为每一个特定的旋转索引都配备一个唯一的旋转密钥。

随着 FHE 程序的规模扩大，所需的旋转索引数量也会随之增长，这往往需要数百 GB 甚至数 TB 的内存来同时存储所有密钥。现有的 FHE 编译器通常将旋转密钥视为一个单一的、静态的实体，在程序执行期间将所有密钥加载到内存中。这种方法缺乏内存可扩展性。相反，一些试图减少内存使用的现有策略（例如 Fhelipe）依赖于将任意旋转索引分解为 2 的幂次链。虽然这减少了存储的密钥数量，但由于链式操作带来的开销，显著增加了计算延迟。目前缺乏能够提供细粒度、自动密钥管理以平衡内存使用和执行速度的编译器。

方法论

作者提出了 KeyMemRT，这是一个旨在单独管理旋转密钥生命周期的编译器和运行时框架。该系统构建在 MLIR（多级中间表示） 基础设施之上，并引入了一种全新的方言（dialect）和运行时架构。

1. `kmrt` MLIR 方言

KeyMemRT 引入了一个名为 kmrt 的新方言，用于在 IR 中直接表示密钥管理。

新类型： 它定义了一种由旋转索引参数化的 rot_key 类型，使编译器能够区分不同索引的密钥。
新操作： 该方言包含了用于显式密钥管理的各种操作：
- kmrt.load_key：加载特定的旋转密钥。
- kmrt.clear_key：从内存中移除一个密钥。
- kmrt.prefetch_key：为运行时提供异步加载的提示。
- kmrt.use_key / kmrt.assume_key：管理密钥的作用域和活跃度（liveness）。
集成： 编译器将高层操作（例如通过 HEIR 从 ckks 方言转换而来）降低（lower）到 openfhe 方言，并在过程中插入 kmrt 操作以管理密钥的生命周期。

2. 编译器优化流水线

KeyMemRT 编译器执行一系列降低和优化过程，以在不牺牲性能的前提下减少内存占用：

CKKS 到 OpenFHE 的降低： 将每一次旋转操作与一个特定的 rot_key 值相关联，将问题从创建活跃范围转变为合并这些范围。
婴儿步-巨步（BSGS）分解： 通过链接旋转操作，将连续旋转索引所需的密钥数量从 $N$ 减少到 $\approx 2\sqrt{N}$ 。这对于矩阵-向量乘法模式特别有效。
合并旋转密钥： 使用数据流分析来检测那些“先清除密钥并立即使用相同索引重新加载”的序列，从而消除冗余的加载/清除对。
旋转提升（Rotation Hoisting）： 识别应用于公共输入的旋转模式，并将其转换为 OpenFHE 库支持的更快的替代方案。
自举密钥管理： 将自举操作（用于刷新密文层级）所使用的内部旋转密钥实例化到 IR 中。这使得编译器能够像管理显式旋转密钥一样，对自举密钥进行细粒度的控制，这是此前未曾解决的功能。
公共子表达式消除（CSE）： 消除由密钥优化产生的死代码。
提示放置（Hint Placement）： 插入 prefetch_key 提示，使运行时能够将密钥加载与计算进行重叠。

3. KeyMemRT 运行时

运行时作为编译后的 C++ 代码与 FHE 库（OpenFHE）之间的中间层。它根据编译器插入的提示，动态管理从磁盘到 RAM 的密钥加载与卸载。

低内存模式： 顺序执行指令，仅在严格需要时才加载密钥，以最小化峰值内存使用量。
平衡模式： 利用预取机制在计算的同时并行加载密钥。它维护一个内部队列，根据用户定义的限制提前加载密钥，并在空间不足时清除密钥，从而降低延迟开销。

核心贡献

新 MLIR 方言 (kmrt)： 第一个将旋转密钥作为具有特定类型和内存管理操作的 IR 元素暴露出来的方言，实现了细粒度分析。
自动细粒度密钥管理： 第一个能够自动追踪密钥活跃度、优化自举密钥，并以单个而非全局集合方式管理密钥生命周期的 FHE 编译器。
运行时框架： 支持按需加载和清除密钥的运行时，利用编译器提示来实现 I/O 与计算的重叠。
集成性： 一个可以被现有 FHE 前端（已通过 Orion 和 HEIR 演示）所针对的后优化编译器工具链。

结果

作者将 KeyMemRT 与两种最先进的方法进行了对比评估：ANT-ACE（将所有密钥保留在内存中）和 Fhelipe（使用分解技术减少密钥但增加了延迟）。基准测试包括各种机器学习模型（MLP, LoLa, LeNet, ResNet 变体, AlexNet），其 FHE 参数设置为环维度 $2^{16}$ 且乘法深度为 30。

内存减少：
- 与 ANT-ACE 相比，KeyMemRT 实现了 1.74x 的几何平均内存减少。在 LeNet 等特定案例中，内存使用量减少了 2.59x。
- 与 Fhelipe 相比，KeyMemRT 实现了 1.16x 的内存减少。
延迟/加速比：
- 与 ANT-ACE 相比，KeyMemRT（平衡模式）实现了 1.20x 的加速。
- 与 Fhelipe 相比，KeyMemRT 实现了 1.73x 的显著加速（几何平均值），在 AlexNet 上提升最高达 2.5x。这是因为 KeyMemRT 避免了 Fhelipe 所需的旋转链式操作带来的计算开销。
可扩展性： KeyMemRT 展示了运行大型程序（如 ResNet-50）的能力，这类程序在 ANT-ACE 下的内存需求超过了 100 GB，而 KeyMemRT 通过保持较低的内存使用量成功运行了这些程序。

重要性与主张

论文声称 KeyMemRT 是第一个提供自动、细粒度密钥管理及自举密钥管理的 FHE 编译器。作者断言，这项工作解锁了 FHE 的内存可扩展性，使得大型程序能够在无需同时存储所有旋转密钥这一沉重负担的情况下运行。

通过将内存使用量与总密钥数解耦，转而基于密钥的活跃度进行管理，KeyMemRT 使 FHE 程序能够在现有硬件上更好地扩展。作者强调，他们的方法并未牺牲性能；不同于通过牺牲速度来换取内存的分解法，KeyMemRT 通过仅在精确需要时加载必要的密钥，同时实现了内存减少和延迟改进。这解决了 FHE 在资源受限或大规模环境中的应用障碍。