TFHE-rs v1.1: 精细化的 GPU 控制和更多算子
本月，Zama 发布了 TFHE-rs v1.1，为 GPU 和 CPU 后端带来了多项重大改进和新功能。
在 GPU 方面，后端已升级为采用与 CPU 相同的默认加密参数，从而将计算错误的可能性降低到 2⁻¹²⁸ 以下——所有这些都对性能影响极小。多 GPU 支持也得到了显著改善：用户现在可以显式选择在哪个 GPU 上运行，从而在 8×H100 GPU 上实现接近每秒 500 次加密 64 位加法。
在 CPU 方面，此版本通过支持更多标量情况来扩展算子集，从而使同态计算更加通用和高效。
在这篇博文中，我们将深入探讨 TFHE-rs v1.1 中的新功能。
更好的多 GPU 吞吐量
在 v1.1 之前，TFHE-rs 的高级 API 使用硬编码策略自动在所有可用 GPU 上调度工作负载，以进行任何加密操作。虽然这对于非常大的整数精度（>128 位）以及大量加载 GPU 的操作（例如乘法）有效，但对于较小的操作（通常是 64 位加密加法或比较）来说并不理想。
从 v1.1 开始，开发人员可以选择对每个操作使用哪个 GPU，从而优化多 GPU 设置的性能。
这是一个快速示例，展示了如何在每个 GPU 上并行执行一百个 64 位加密加法，其中每个加法都在单个 GPU 上计算：
use tfhe::{ConfigBuilder, set_server_key, ClientKey, CompressedServerKey, FheUint64, GpuIndex};
use tfhe::prelude::*;
use rayon::prelude::*;
use tfhe::core_crypto::gpu::get_number_of_gpus;
use rand::{thread_rng, Rng};
fn main() {
let config = ConfigBuilder::default().build();
let client_key = ClientKey::generate(config);
let compressed_server_key = CompressedServerKey::new(&client_key);
let num_gpus = get_number_of_gpus();
let sks_vec = (0..num_gpus)
.map(|i| compressed_server_key.decompress_to_specific_gpu(GpuIndex::new(i)))
.collect::>();
let batch_size = num_gpus * 100;
let mut rng = thread_rng();
let left_inputs = (0..batch_size)
.map(|_| FheUint64::encrypt(rng.gen::(), &client_key))
.collect::>();
let right_inputs = (0..batch_size)
.map(|_| FheUint64::encrypt(rng.gen::(), &client_key))
.collect::>();
let chunk_size = (batch_size / num_gpus) as usize;
left_inputs
.par_chunks(chunk_size)
.zip(
right_inputs
.par_chunks(chunk_size)
)
.enumerate()
.for_each(
|(i, (left_inputs_on_gpu_i, right_inputs_on_gpu_i))| {
left_inputs_on_gpu_i
.par_iter()
.zip(right_inputs_on_gpu_i.par_iter())
.for_each(|(left_input, right_input)| {
set_server_key(sks_vec[i].clone());
left_input + right_input;
});
},
);
}
这里发生了什么？与使用 TFHE-rs 进行的常规 GPU 计算不同的是，服务器密钥的定义方式：
let sks_vec = (0..num_gpus)
.map(|i| compressed_server_key.decompress_to_specific_gpu(GpuIndex::new(i)))
.collect::>();
在这里，创建了一个服务器密钥向量，每个密钥都在特定的 GPU 上。
然后，不是像通常那样在输入上调用 par_iter()，而是将输入分块以分布到所有 GPU 上：
left_inputs
.par_chunks(chunk_size)
.zip(
right_inputs
.par_chunks(chunk_size)
)
.enumerate()
.for_each(
|(i, (left_inputs_on_gpu_i, right_inputs_on_gpu_i))| {
left_inputs_on_gpu_i
.par_iter()
.zip(right_inputs_on_gpu_i.par_iter())
.for_each(|(left_input, right_input)| {
set_server_key(sks_vec[i].clone());
left_input + right_input;
});
},
);
通过使用 set_server_key(sks_vec[i].clone()) 设置与每个块关联的 GPU 对应的服务器密钥，加法将在所有 GPU 上独立计算。请注意，当执行 sks_vec[i].clone() 时，服务器密钥的指针会复制到特定线程，而不是服务器密钥本身的内容，因此不会产生额外的开销。你可以按照我们的专用教程进一步最大化多 GPU 吞吐量。
通过设置此逻辑，TFHE-rs 现在可以在 8xH100 GPU 上实现接近每秒 500 次 64 位加密整数的加法。
CPU 后端上的新算子
TFHE-rs v1.1 为 CPU 后端带来了一些新增功能和改进：
-
对 select 的标量支持：以前，select 操作仅适用于加密值。在 v1.1 中，你现在可以使用标量（明文）值作为可选择的操作数。对于 64 位输入，此操作大约在 20 毫秒内执行。
-
改进的减法：减法现在支持左侧的标量，从而可以进行像 scalar - encrypted 这样的表达式。对于 64 位操作数，此操作大约需要 79 毫秒。
-
新的点积算子：v1.1 引入了 FheBool 值的向量与任何支持的标量类型之间的点积运算。在包含 1,024 个元素的向量上，执行时间约为 2 秒。
所有性能基准均在 AWS hpc7a.96xlarge 实例上测量。
更智能的密钥生成
为了更好地支持内存受限环境中的操作，v1.1 还引入了“分块”引导密钥生成。此功能允许以较小的块生成引导密钥，这些块稍后可以在用于加密计算的更高容量的服务器上组装成完整密钥。
TFHE-rs 的下一个版本将继续提高性能并引入新功能。敬请关注！
额外链接
