基于秘密共享重构 DeepSeek DeepGEMM Kernel 的安全高效 MPC-GEMM 方案

发布于 2025-3-11 02:10

浏览

0收藏

摘要

本文针对安全多方计算（MPC）框架下通用矩阵乘法（GEMM）运算的性能瓶颈，提出一种全新的 MPC-GEMM 实现方案。该方案的核心思想在于：基于加法秘密共享重构 DeepSeek DeepGEMM 的 CUDA kernel，将 MPC 协议的逻辑与 DeepGEMM 的底层优化深度融合，消除 MPC 协议与 GPU 计算之间的“两张皮”现象。方案采用 INT8/FP8 数据表示、秘密共享运算的 kernel 级实现、Hopper 架构优化（如适用）、GPU 并行 Beaver 三元组生成以及 JIT 编译等关键技术。本文将详细阐述方案的设计原理、技术框架、实现细节（包括 kernel 代码示例、算法描述、优化策略），并从可行性、安全性、高效性等方面进行全面深入的论证，最后与其他 MPC-GEMM 方案进行对比。方案旨在实现真正意义上的安全、高效的 MPC-GEMM，为隐私保护机器学习提供强有力的支持。

关键词： DeepGEMM, DeepSeek, MPC, GEMM, 秘密共享, CUDA, Kernel 重构, 安全计算, INT8, FP8, Hopper 架构, Beaver 三元组, JIT 编译, 并行计算

1. 引言：MPC-GEMM 的性能挑战与 DeepGEMM 的机遇

安全多方计算（MPC）使得互不信任的参与方能够在不泄露各自私有数据的前提下进行协同计算，是实现隐私保护机器学习的关键技术。通用矩阵乘法（GEMM）作为深度学习模型的核心运算，其在 MPC 框架下的实现（MPC-GEMM）的效率直接影响着隐私保护机器学习应用的整体性能和实用性。然而，现有的 MPC-GEMM 方案普遍面临着严重的性能挑战：

计算开销：MPC 协议的密码学运算（如秘密共享、同态加密）计算复杂度远高于明文计算。
通信开销：多数 MPC 协议需要在参与方之间进行大量的交互通信，尤其是在执行乘法运算时，通信开销成为主要瓶颈。
硬件加速：如何在 MPC 的安全约束下有效利用 GPU 等硬件加速器进行计算，是一个极具挑战性的问题。

传统的 MPC-GEMM 方案通常采用“两张皮”模式：MPC 协议负责保证计算的安全性，GPU 负责提供计算加速，两者之间通过某种安全接口（如可信执行环境 TEE 或同态加密）进行交互。这种模式的缺点在于：

交互开销：MPC 协议与 GPU 计算之间存在数据转换（如明文与密文、秘密份额与 GPU 可处理格式之间的转换）和通信的开销，限制了整体性能。
GPU 利用率：GPU 计算部分通常受到 MPC 协议的制约，无法充分发挥 GPU 的并行计算能力和 DeepGEMM 等底层优化库的性能优势。

DeepSeek 最新发布的 DeepGEMM 是一个为 NVIDIA GPU 优化的高性能 GEMM 库。它通过 FP8 低精度计算、针对 GPU 架构的优化、CUDA kernel 优化以及 JIT 编译等技术，大幅提升了 GEMM 运算的效率。虽然 DeepGEMM 并非专门为 MPC 设计，但其在 kernel 级别的优化为我们提供了一个重要的机遇：能否将 MPC 协议与 DeepGEMM 的底层优化进行深度融合，消除“两张皮”现象，实现真正意义上的安全高效的 MPC-GEMM？

2. 方案原理：深度融合 MPC 与 DeepGEMM

基于 MPC 与 DeepGEMM 的深度融合，就可以尝试构想一种全新的 MPC-GEMM 方案：基于秘密共享重构 DeepSeek DeepGEMM kernel。该方案的核心思想是：将 MPC 协议中与 GEMM 运算相关的计算逻辑（秘密份额的加法、乘法）直接实现在 DeepGEMM 的 CUDA kernel 中，让 GPU 直接执行一个完整的“MPC-GEMM”运算。

方案的设计基于以下几个关键原理：

1）加法秘密共享：采用加法秘密共享作为 MPC 的基础安全机制。加法秘密共享具有以下优点：

简单高效：实现简单，只需要进行模加运算。
加法同态：秘密份额的加法对应于明文的加法，使得加法运算可以在本地高效执行，无需通信。
安全性：信息论安全，只要参与方不合谋，任何单独的秘密份额都不会泄露关于原始数据的任何信息。

2）INT8/FP8 数据表示：为了降低计算和通信开销，我们借鉴 DeepGEMM 对低精度计算的使用，将输入数据（FP32/FP64/定点数）映射到 INT8 或 FP8。

INT8 映射：对于 INT8，我们采用偏移映射等策略，充分利用 INT8 的表示范围，并简化秘密共享运算。
FP8 映射:如果采用 FP8，可以利用 DeepGEMM 自身的 FP8 支持。

3）DeepGEMM Kernel 重构：方案的核心在于对 DeepGEMM 的 CUDA kernel 进行重构。我们将 MPC 协议的逻辑（即秘密共享下的加法和乘法）直接嵌入到 kernel 中。

输入/输出：Kernel 的输入和输出直接是秘密份额（INT8 或 FP8），而不是明文数据。
基本运算：将 kernel 中的加法和乘法替换为 MPC 协议下的秘密共享加法和乘法（基于 Beaver 三元组）。
保留优化：尽最大可能保留 DeepGEMM 原有的针对 GPU 架构的优化技术，如 tiling、loop unrolling、shared memory 利用、warp-level primitives、指令级并行等，并针对秘密共享运算进行适配。
异步计算: 尽可能利用GPU的异步计算能力。

4）Beaver 三元组乘法：为了在秘密共享下实现乘法，采用 Beaver 三元组乘法协议。可以在 kernel 中实现 Beaver 三元组乘法协议，并利用 warp-level primitives（如__shfl_xor_sync）进行优化。

5）GPU 并行 Beaver 三元组生成：为了提高 Beaver 三元组的生成效率，并减少预处理阶段的通信开销，我们可以利用 GPU 的并行计算能力，在 GPU 上并行生成 Beaver 三元组。

6）JIT 编译：我们充分利用 DeepGEMM 的 JIT 编译技术（如果 DeepGEMM 提供 JIT 编译接口；如果没有，我们可以自行实现 JIT 编译），根据 GEMM 形状、块大小、参与方数量等参数，动态生成高度优化的 MPC-GEMM kernel。

7）简化的 MPC 协议：由于 GPU 直接参与 MPC 协议的执行（我们将其视为一个“半诚实”的参与方），我们可以简化 MPC 协议的设计，减少通信轮数和通信量。

3. 技术框架与实现细节

3.1 技术框架

方案的技术框架主要由以下几个模块构成：

秘密共享模块：

a.负责将参与方的输入数据（FP32、FP64 或定点数）进行加法秘密共享。

b.将秘密份额转换为 INT8 或 FP8 表示（通过映射）。

c.实现秘密共享上的加法和乘法运算（基于 Beaver 三元组）。

d.提供秘密份额的生成、分发、重构等功能。

DeepGEMM Kernel 重构模块：

a.负责对 DeepGEMM 的 CUDA kernel 进行重构，将秘密共享运算（加法和乘法）嵌入到 kernel 中。

b.保留并适配 DeepGEMM 原有的 GPU 架构优化。

c.利用 JIT 编译技术（或手动实现），动态生成针对特定参数（GEMM 形状、块大小、参与方数量等）的优化 kernel。

MPC 协议协调模块：

a.负责协调各参与方和 GPU 之间的交互。

b.管理 Beaver 三元组的分发（如果采用离线生成）。

c.触发 GPU kernel 的执行。

GPU Beaver 三元组生成模块：

a.利用 GPU 的并行计算能力，高效生成 Beaver 三元组。

3.2 工作流程

整个 MPC-GEMM 的计算流程分为离线阶段和在线阶段：

离线阶段（预处理）：

利用 GPU 并行生成 Beaver 三元组，并将三元组的秘密份额分发给各参与方（和 GPU 线程）。

在线阶段：

参与方收集各自的输出份额。
将对应位置的份额相加（模运算，如果是 INT8；浮点加法，如果是 FP8），重构出最终的 GEMM 结果。
如果需要，可以将结果转换回 FP32 或 FP64 格式。
Kernel 计算完成后，输出结果仍然是秘密份额（INT8 或 FP8）的形式。
GPU 将输出份额返回给参与方。
GPU 执行重构后的 DeepGEMM kernel。
在 kernel 内部：
整个计算过程高度并行化。
将输入数据（秘密份额）和 Beaver 三元组份额加载到 shared memory。
使用 tiling 技术将矩阵分块。
对于每个块，执行秘密共享下的加法和乘法运算（利用 Beaver 三元组和 warp-level primitives）。
利用 GPU 架构优化（如 tiling, loop unrolling, shared memory, warp-level primitives, 指令级并行, 异步计算等）。
将中间结果累加到 shared memory 或 registers 中。
MPC 协议协调模块根据 GEMM 运算的参数（形状、块大小等）和参与方数量，触发 DeepGEMM Kernel 重构模块生成相应的 CUDA kernel（利用 JIT 编译或手动实现）。
参与方将各自持有的秘密份额（INT8 或 FP8）直接作为输入，传递给生成的 CUDA kernel。
每个参与方将自己的输入矩阵的每个元素进行加法秘密共享。
将秘密份额转换为 INT8 或 FP8 表示。
输入准备：
Kernel 调用：
GPU 并行计算：
输出处理：
结果重构：

3.3 关键实现细节

本文中所用代码均是伪代码，根据通义灵码的建议生成的，只能看出大致的意思，不能直接使用。

3.3.1 数据表示

INT8 映射 (如果采用 INT8)：
我们推荐使用偏移映射。假设原始数据为 FP32，映射规则如下：
映射公式：
其中，S1、S2 是缩放因子，O1、O2 是偏移量。具体数值需要根据实际数据分布和 INT8 的表示范围来确定。
对于 FP32 正数 x：INT8 = round(x * S1) + O1
对于 FP32 负数 x：INT8 = round(x * S2) + O2
将 FP32 的 NaN 映射到 INT8 的 -128。
将 FP32 的 +Inf 映射到 INT8 的 -127。
将 FP32 的 -Inf 映射到 INT8 的 -126。
将 FP32 的 0 映射到 INT8 的 0。
将 FP32 的其他正数，等比例映射到 INT8 的 [1, 127] 区间。
将 FP32 的其他负数，等比例映射到 INT8 的 [-125, -1] 区间。
FP8 表示 (如果采用 FP8):如果采用FP8，可以直接利用DeepGEMM对FP8的支持。

3.3.2 CUDA Kernel 中的秘密共享乘法

以下是 CUDA kernel 中实现秘密共享乘法（基于加法秘密共享和 Beaver 三元组）的示例代码，并加入了详细注释：

#include <cooperative_groups.h>

namespace cg = cooperative_groups;

template<typename T>
__global__ void mpc_gemm_kernel(T* x_shares, T* y_shares,
                                  T* a_shares, T* b_shares, T* c_shares,
                                  T* z_shares,
                                  int m, int n, int k, int num_parties) {
// 获取线程 ID、块 ID 以及块维度
int tid = threadIdx.x;
int bid_x = blockIdx.x;
int bid_y = blockIdx.y;
int block_dim = blockDim.x;

// 定义 shared memory 变量 (使用双缓冲)
  __shared__ T x_shared[2][BLOCK_SIZE][BLOCK_SIZE];
  __shared__ T y_shared[2][BLOCK_SIZE][BLOCK_SIZE];
  __shared__ T a_shared[2][BLOCK_SIZE][BLOCK_SIZE];
  __shared__ T b_shared[2][BLOCK_SIZE][BLOCK_SIZE];
  __shared__ T c_shared[2][BLOCK_SIZE][BLOCK_SIZE];

// 使用 cooperative groups
  cg::thread_block cta = cg::this_thread_block();
  cg::grid_group grid = cg::this_grid();
  cg::thread_block_tile<32> warp = cg::tiled_partition<32>(cta);

// 计算当前线程负责的矩阵元素的坐标
int row = bid_y * BLOCK_SIZE + tid / BLOCK_SIZE;
int col = bid_x * BLOCK_SIZE + tid % BLOCK_SIZE;

// 初始化累加器
  T acc = 0;

// 循环处理矩阵块 (tiling)
    int buffer_idx = 0; // 双缓冲索引
    for (int i = 0; i < k; i += BLOCK_SIZE) {
      // 将数据从全局内存加载到 shared memory (异步加载, 如果支持)
       if (grid.rank() == 0 && i + BLOCK_SIZE < k) {
            //仅rank 0 的block进行异步加载
            //这里只是伪代码，实际使用需要根据数据类型进行调整
            cudaMemcpyAsync(&x_shared[(buffer_idx+1)%2][0][0], &x_shares[(row * k) + i + BLOCK_SIZE], BLOCK_SIZE * BLOCK_SIZE * sizeof(T), cudaMemcpyDeviceToDevice);
            cudaMemcpyAsync(&y_shared[(buffer_idx+1)%2][0][0], &y_shares[((i + BLOCK_SIZE) * n) + col], BLOCK_SIZE * BLOCK_SIZE * sizeof(T), cudaMemcpyDeviceToDevice);
            cudaMemcpyAsync(&a_shared[(buffer_idx+1)%2][0][0], &a_shares[(row * k) + i + BLOCK_SIZE], BLOCK_SIZE * BLOCK_SIZE * sizeof(T), cudaMemcpyDeviceToDevice);
            cudaMemcpyAsync(&b_shared[(buffer_idx+1)%2][0][0], &b_shares[((i + BLOCK_SIZE) * n) + col], BLOCK_SIZE * BLOCK_SIZE * sizeof(T), cudaMemcpyDeviceToDevice);
            cudaMemcpyAsync(&c_shared[(buffer_idx+1)%2][0][0], &c_shares[row*n + col], BLOCK_SIZE*BLOCK_SIZE*sizeof(T), cudaMemcpyDeviceToDevice);
        }

        if(row < m && (i + tid % BLOCK_SIZE) < k){
          x_shared[buffer_idx][tid / BLOCK_SIZE][tid % BLOCK_SIZE] = x_shares[row * k + (i + tid % BLOCK_SIZE)];
        } else {
          x_shared[buffer_idx][tid / BLOCK_SIZE][tid % BLOCK_SIZE] = 0;
        }

        if((i + tid / BLOCK_SIZE) < k && col < n) {
          y_shared[buffer_idx][tid / BLOCK_SIZE][tid % BLOCK_SIZE] = y_shares[(i + tid / BLOCK_SIZE) * n + col];
        } else {
           y_shared[buffer_idx][tid / BLOCK_SIZE][tid % BLOCK_SIZE] = 0;
        }
      
        if(row < m && (i + tid % BLOCK_SIZE) < k){
            a_shared[buffer_idx][tid / BLOCK_SIZE][tid % BLOCK_SIZE] = a_shares[row*k + (i + tid%BLOCK_SIZE)];
        } else {
            a_shared[buffer_idx][tid / BLOCK_SIZE][tid % BLOCK_SIZE] = 0;
        }

        if((i + tid / BLOCK_SIZE) < k && col < n){
          b_shared[buffer_idx][tid / BLOCK_SIZE][tid % BLOCK_SIZE] = b_shares[(i + tid / BLOCK_SIZE)*n + col];
        } else {
           b_shared[buffer_idx][tid / BLOCK_SIZE][tid % BLOCK_SIZE] = 0;
        }
      
        ```cuda
        if(row < m && col < n) {
          c_shared[buffer_idx][tid / BLOCK_SIZE][tid % BLOCK_SIZE] = c_shares[row*n + col];
        } else {
           c_shared[buffer_idx][tid / BLOCK_SIZE][tid % BLOCK_SIZE] = 0;
        }

      cta.sync(); // 等待所有线程加载完成, 以及异步加载完成

      // 计算当前块的乘积 (循环展开)
      #pragma unroll
      for (int j = 0; j < BLOCK_SIZE; ++j) {
        // 计算 d = x - a 和 e = y - b (本地计算)
        T d_local = x_shared[buffer_idx][tid / BLOCK_SIZE][j] - a_shared[buffer_idx][tid / BLOCK_SIZE][j];
        T e_local = y_shared[buffer_idx][j][tid % BLOCK_SIZE] - b_shared[buffer_idx][j][tid % BLOCK_SIZE];

        // 使用 warp-level shuffle 指令计算 d 和 e 的全局和
        T d_global = 0;
        T e_global = 0;

        #pragma unroll
        for (int w = 0; w < warp.size(); ++w) {
          d_global += warp.shfl_xor(d_local, w);
          e_global += warp.shfl_xor(e_local, w);
        }

        // 计算 z = c + d * b + e * a + d * e (本地计算)
        // 手动进行指令级并行
        T term1 = d_local * b_shared[buffer_idx][j][tid % BLOCK_SIZE];
        T term2 = e_local * a_shared[buffer_idx][tid / BLOCK_SIZE][j];
        T term3 = d_global * e_global;
          
        //根据数据类型进行模运算
        if constexpr (std::is_same_v<T, int8_t>) {
            acc = (acc + c_shared[buffer_idx][tid / BLOCK_SIZE][tid % BLOCK_SIZE] + term1 + term2 + term3) & 0xFF;
        } else {
             acc += c_shared[buffer_idx][tid / BLOCK_SIZE][tid % BLOCK_SIZE] + term1 + term2 + term3;
        }
          
      }

      cta.sync(); // 确保所有线程完成当前块的计算
      buffer_idx = (buffer_idx + 1) % 2;
    }
    // 将结果写回全局内存
    if(row < m && col < n){
        z_shares[row * n + col] = acc;
    }
}

代码解释：

模板参数 T：使用模板参数 T，可以支持 INT8 和 FP8 两种数据类型。
双缓冲 (Double Buffering)：使用两组 shared memory 数组，实现计算和数据加载的流水线操作。
异步数据加载：在外层循环的开始处，尝试使用 cudaMemcpyAsync 异步地将下一批次的数据从全局内存加载到 shared memory。
Cooperative Groups：使用 Cooperative Groups 提供的 thread_block、grid_group 和 thread_block_tile 类型来更精细地控制线程块和 warp 级别的并行。
Warp-level Shuffle 指令优化：

a.使用 warp.shfl_xor(val, lane) 替代 __shfl_xor_sync(mask, val, lane)。

b.循环展开 warp-level shuffle 操作。

指令级并行（手动）：在计算 z 时，将乘法和加法运算交错进行，尽可能利用 GPU 的指令级并行能力。
循环展开：使用 #pragma unroll 指令展开内层循环。
模运算: 如果 T 是 int8_t，则使用 & 0xFF 进行模 256 运算。
Tiling: 使用tiling技术将矩阵分块处理。
并行化:

a.线程块 (Block)：不同的线程块负责计算输出矩阵 Z 的不同块（tiling）。

b.线程 (Thread)：线程块内部的线程协同计算秘密共享乘法。

3.3.3 Hopper 架构优化（深化）

TMA (Tensor Memory Accelerator)：通过流水线、双缓冲和 cudaMemcpyAsync，尽可能利用 TMA 的异步数据传输能力，隐藏内存访问延迟。
Tensor Core 利用：

#include <mma.h>
usingnamespace nvcuda;
// ...
 wmma::fragment<wmma::matrix_a, 16, 16, 16, int8_t, wmma::row_major> a_frag;
 wmma::fragment<wmma::matrix_b, 16, 16, 16, int8_t, wmma::col_major> b_frag;
 wmma::fragment<wmma::accumulator, 16, 16, 16, int32_t> c_frag;
 wmma::fragment<wmma::accumulator, 16, 16, 16, int32_t> acc_frag;

 wmma::fill_fragment(acc_frag, 0);
for (int i = 0; i < k; i += 16) {
     wmma::load_matrix_sync(a_frag, &x_shared[...], ...); // 加载数据到 fragment, 需要根据实际情况填写参数
     wmma::load_matrix_sync(b_frag, &y_shared[...], ...); // 加载数据到 fragment, 需要根据实际情况填写参数

     wmma::mma_sync(acc_frag, a_frag, b_frag, acc_frag); // 矩阵乘累加
 }
//将秘密共享乘法结果加到acc_frag上
 wmma::store_matrix_sync(&c_shared[...], acc_frag, ... , wmma::mem_row_major); // 存储结果
 ```
*   **FP8 计算:** 如果采用FP8, 可以直接使用DeepGEMM中针对FP8和Tensor Core的优化。
*   **数据类型转换：** 如果 INT8 的 `wmma.mma` 指令效果不佳，可考虑将 INT8 份额转换为 FP16 或 INT32，然后使用相应的 `wmma.mma` 指令。但类型转换也需在秘密共享下进行。

a.INT8 计算：尝试使用 wmma::mma_s8s8s32 指令进行 INT8 矩阵乘法：

Shared Memory 优化：

a.通过 tiling 技术和合理的数据访问模式，最大程度地复用 shared memory 中的数据。

b.合理安排 shared memory 中数据的存储位置，避免 bank conflict。

Warp-level Primitives 与指令级并行：

a.充分利用__shfl_xor_sync 或warp.shfl_xor 指令在 warp 内部高效地进行数据交换和规约求和。

b.在 kernel 代码中，尽可能地将独立的指令放在一起执行，利用 GPU 的指令级并行能力。

3.3.4 GPU 并行 Beaver 三元组生成

算法：

初始化 cuRAND：在每个线程中初始化一个 cuRAND 伪随机数生成器状态。
生成随机数：使用 cuRAND 库在每个线程中并行生成三个 INT8 或 FP8 类型的随机数（a, b, c）。
验证三元组：在每个线程中验证生成的三元组是否满足 Beaver 三元组的条件（c == a * b）。
秘密共享：在 kernel 中直接对验证通过的三元组 (a, b, c) 进行加法秘密共享。
存储份额：将每个参与方的三元组份额存储到全局内存中的一个数组中。

CUDA Kernel 代码示例（INT8）：

#include <curand_kernel.h>

struct BeaverTripleShares {
    int8 a_share;
    int8 b_share;
    int8 c_share;
};

__global__ void generate_beaver_triples(BeaverTripleShares* triples, int num_triples, int num_parties) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;

    // 初始化 cuRAND 伪随机数生成器
    curandState_t state;
    curand_init(blockIdx.x * blockDim.x + threadIdx.x, 0, 0, &state);

    // 生成 Beaver 三元组并进行秘密共享
    if (tid < num_triples) {
        // 1. 生成随机数
        int8 a = (int8)curand(&state);
        int8 b = (int8)curand(&state);
        int8 c = (int8)curand(&state);

        // 2. 验证三元组 (注意处理溢出)
        if (((int)a * (int)b & 0xFF) == (c & 0xFF)) {
            // 3. 进行秘密共享
            int8 a_shares[num_parties];
            int8 b_shares[num_parties];
            int8 c_shares[num_parties];

            int8 a_sum = 0;
            int8 b_sum = 0;
            int8 c_sum = 0;

            for (int i = 0; i < num_parties - 1; i++) {
                a_shares[i] = (int8)curand(&state);
                b_shares[i] = (int8)curand(&state);
                c_shares[i] = (int8)curand(&state);

                a_sum += a_shares[i];
                b_sum += b_shares[i];
                c_sum += c_shares[i];
            }

            a_shares[num_parties - 1] = a - a_sum; // 加法秘密共享
            b_shares[num_parties - 1] = b - b_sum;
            c_shares[num_parties - 1] = c - c_sum;
            
             a_shares[num_parties - 1] = a_shares[num_parties-1] & 0xFF;
             b_shares[num_parties - 1] = b_shares[num_parties-1] & 0xFF;
             c_shares[num_parties - 1] = c_shares[num_parties-1] & 0xFF;

            // 4. 存储秘密份额
            for (int i = 0; i < num_parties; i++) {
                triples[tid * num_parties + i].a_share = a_shares[i];
                triples[tid * num_parties + i].b_share = b_shares[i];
                triples[tid * num_parties + i].c_share = c_shares[i];
            }
        } else {
           // 如果验证失败，可以将其设置为一个特殊值（如全 0），
           for (int i = 0; i < num_parties; i++) {
                triples[tid * num_parties + i].a_share = 0;
                triples[tid * num_parties + i].b_share = 0;
                triples[tid * num_parties + i].c_share = 0;
            }
        }
    }
}

代码解释：

curand_kernel.h：包含了 cuRAND 库的函数声明。
BeaverTripleShares 结构体：定义了 Beaver 三元组份额的结构。
generate_beaver_triples kernel：

a.在 kernel 中直接对验证通过的 Beaver 三元组 (a, b, c) 进行加法秘密共享。

b.为每个参与方生成随机份额。

c.最后一个参与方的份额通过总和与其他份额的差值计算得到。

((int)a * (int)b & 0xFF)：计算 a * b (mod 256)。
(c & 0xFF)：取 c 的低 8 位。
triples：指向全局内存中存储 Beaver 三元组份额的数组的指针，其大小应为 num_triples * num_parties。
num_triples：要生成的 Beaver 三元组的数量。
num_parties：参与方的数量。
tid：线程 ID。
curandState_t：cuRAND 伪随机数生成器的状态。每个线程都需要一个独立的状态。
curand_init：初始化伪随机数生成器。这里使用线程 ID 作为种子，确保每个线程生成的随机数序列不同。
curand：生成一个 32 位无符号整数随机数。我们将其强制转换为 int8。
验证三元组：
秘密共享：
存储份额：将每个参与方的三元组份额存储到 triples 数组中。

使用方法：

在 GPU 上分配足够大的内存来存储 Beaver 三元组的所有份额 (BeaverTripleShares* triples)。
调用 generate_beaver_triples kernel，生成 Beaver 三元组并进行秘密共享。
在 MPC-GEMM kernel 中，每个线程根据其线程 ID 和参与方 ID 从 triples 数组中获取相应的 Beaver 三元组份额。

优化：

可以通过增加线程块和线程数量来进一步提高 Beaver 三元组生成的并行度。
可以使用更高效的随机数生成器（如 Philox 算法）来提高随机数生成的速度和质量。
可以将 Beaver 三元组的生成、验证和秘密共享融合到一个 kernel 中，减少数据传输开销。

3.3.5 JIT 编译优化

JIT 编译技术允许我们在运行时根据具体的参数动态生成优化的 CUDA kernel 代码。在 MPC-GEMM 中，我们可以利用 JIT 编译进行以下优化：

代码特化：

GEMM 参数：根据 GEMM 运算的形状（M, N, K）、块大小（BLOCK_SIZE）、数据类型（INT8 或 FP8）等参数，生成专门针对这些参数优化的 kernel 代码。例如，可以根据 M、N、K 的大小选择最合适的 tiling 策略和 shared memory 使用方式。
MPC 参数：根据参与方数量、秘密共享方案（加法秘密共享）等参数，生成相应的 kernel 代码。例如，如果参与方数量较少，可以使用更激进的 warp-level shuffle 优化。
Hopper 架构特性：根据目标 GPU 的计算能力（Compute Capability），启用或禁用某些 Hopper 架构特有的优化（如 TMA）。

常量折叠：

Beaver 三元组内联：如果 Beaver 三元组是在预处理阶段生成的，并且在 kernel 执行期间不会改变，可以将三元组的份额直接作为编译时常量内联到 kernel 代码中，减少运行时内存访问。
其他常量：将参与方数量、块大小、GEMM 形状等参数也作为编译时常量内联到 kernel 代码中，允许编译器进行更多的优化（如常量传播、死代码消除等）。

循环展开：

根据 GEMM 形状和块大小，对 kernel 中的循环进行部分或完全展开，减少循环控制开销，并增加指令级并行度。
特别是对于秘密共享乘法协议中的内层循环，可以进行更激进的展开。

指令级并行：

JIT 编译器可以分析 kernel 代码中的数据依赖关系，尽可能地将独立的指令放在一起执行，利用 GPU 的指令级并行能力。
我们可以手动调整 kernel 代码中的指令顺序，以帮助编译器更好地进行指令级并行优化。

自动调整block size和grid size: 可以根据矩阵规模、数据类型等，自动调整kernel的block size和grid size，以充分利用GPU资源。

实现方式：

NVRTC (NVIDIA Runtime Compilation)：NVRTC 是 NVIDIA 提供的一个运行时编译库，可以在程序运行时将 CUDA C++ 代码编译为 PTX 汇编代码，然后加载到 GPU 中执行。
NVCC (NVIDIA CUDA Compiler)：NVCC 是 NVIDIA 的 CUDA 编译器，也可以用于 JIT 编译。可以在编译时使用 -D 选项定义宏，然后在 kernel 代码中使用 #ifdef 等预处理指令来根据不同的宏定义生成不同的代码。

示例：

假设我们要根据参与方数量 n 进行代码特化。我们可以在 kernel 代码中使用如下预处理指令：

#if N_PARTIES == 2
    // 针对 2 个参与方的优化代码
    int8 d_global = __shfl_xor_sync(0xFFFFFFFF, d_local, 1);
    ```C++
    int8 e_global = __shfl_xor_sync(0xFFFFFFFF, e_local, 1);
#elif N_PARTIES == 3
    // 针对 3 个参与方的优化代码
     int8 d_global = 0;
      int8 e_global = 0;
      for (int w = 0; w < warp.size(); ++w) {
        d_global += warp.shfl_xor(d_local, w);
        e_global += warp.shfl_xor(e_local, w);
      }
#else
    // 通用代码
#endif

在编译时，通过 -D 选项指定 N_PARTIES 的值，NVCC 或 NVRTC 就会生成针对特定参与方数量的优化 kernel 代码。

4. 方案论证

4.1 可行性论证

DeepGEMM Kernel 可修改性：DeepGEMM 的 CUDA kernel 本质上是 C/C++ 代码，可以进行修改和扩展。
秘密共享运算可实现性：加法秘密共享和基于 Beaver 三元组的乘法协议都可以在 INT8 或 FP8 数据类型上高效实现。
GPU 并行计算可行性：CUDA 编程模型支持细粒度的并行计算，可以充分利用 GPU 的并行计算能力。
JIT 编译可行性：JIT 编译技术已经广泛应用，NVRTC 和 NVCC 都提供了 JIT 编译功能。
Hopper 架构优化可行性：Hopper 架构的特性（TMA、Tensor Core、Shared Memory、Warp-level Primitives）都可以在 CUDA 编程中加以利用。

4.2 安全性论证

本方案的安全性基于以下几个方面：

秘密共享的安全性：采用的加法秘密共享方案是信息论安全的，只要参与方不合谋，任何单独的秘密份额都不会泄露关于原始数据的任何信息。
Beaver 三元组乘法协议的安全性：Beaver 三元组乘法协议在半诚实模型下是安全的。只要 Beaver 三元组是独立于输入数据生成的，并且参与方诚实地执行协议，攻击者就无法从公开的中间值（d 和 e）中推断出关于秘密输入（x 和 y）的任何信息。
GPU 计算的安全性：

GPU 始终只接触到秘密份额，无法获得任何关于明文数据的信息。
重构后的 DeepGEMM kernel 只执行秘密共享运算，不包含任何可能泄露敏感信息的操作（如直接访问内存地址、向外部发送数据等）。
即使攻击者控制了 GPU，也只能获得秘密份额，无法恢复出原始数据。

JIT 编译的安全性：

JIT 编译器生成的 kernel 代码只包含必要的秘密共享运算和优化逻辑，不包含任何恶意代码。
可以对 JIT 编译器生成的代码进行静态分析和安全审计。

抵御侧信道攻击：

虽然 GPU 内部的计算对参与方透明，但仍然需要考虑侧信道攻击（如时间攻击、功耗攻击）。
可以采用掩码（masking）技术来防御侧信道攻击。具体来说，可以将秘密份额与一个随机数进行运算（如异或），然后在掩码后的份额上进行计算，最后再去除掩码。
可以对 kernel 代码进行随机化，使得每次执行的指令顺序和内存访问模式都不同，增加侧信道攻击的难度。

4.3 高效性论证

相比于传统的“两张皮”MPC-GEMM 方案，本方案具有以下优势：

消除交互开销：将 MPC 协议逻辑直接嵌入到 DeepGEMM kernel 中，彻底消除了 MPC 协议与 GPU 计算之间的所有交互开销（如数据格式转换、安全通道传输等）。这是本方案相对于传统方案最大的优势所在。
充分利用 DeepGEMM 优化：GPU 直接执行 MPC-GEMM 运算，可以充分利用 DeepGEMM 原有的针对 GPU 架构的各种优化（tiling、loop unrolling、shared memory 利用、TMA、Tensor Core、warp-level primitives、指令级并行等）。
低精度计算：使用 INT8 或 FP8 数据类型，相比于 FP32 或 FP64，可以显著减少计算量和通信量。
GPU 并行 Beaver 三元组生成：利用 GPU 并行生成 Beaver 三元组，大幅减少了预处理阶段的开销。
简化的 MPC 协议：将 GPU 视为“半诚实”参与方，可以简化 MPC 协议的设计，减少通信轮数。
JIT 编译优化：通过 JIT 编译，可以针对具体的 GEMM 参数和 MPC 参数生成高度定制化的 kernel 代码，进一步提升性能。
高度并行化: 秘密共享的加法、乘法，Beaver三元组的生成都可以在GPU上高度并行。

量化分析（举例）：

假设一个 MPC-GEMM 运算涉及两个矩阵 A 和 B 的乘法，矩阵大小为 1024x1024，参与方数量为 3。

传统“两张皮”方案：
整个过程中，数据至少需要在网络上传输 3 次（输入 2 次，输出 1 次），并且涉及到多次数据格式转换。

参与方之间需要通过网络传输秘密份额（FP32 或 FP64）。
需要将秘密份额转换为 GPU 可处理的格式（如加密）。
GPU 执行 GEMM 计算。
将计算结果（加密或编码）传输回参与方。
参与方进行解密或解码，并重构结果。

本方案：
整个过程中，数据只需要在网络上传输 2 次（输入和输出），并且都是 INT8 类型，数据量大大减少。GPU 内部的计算高度优化，且无需与 MPC 协议进行交互。

参与方将输入数据进行秘密共享，并映射到 INT8。
参与方将 INT8 秘密份额直接发送给 GPU（通过 MPI 等）。
GPU 执行重构后的 DeepGEMM kernel，直接在 INT8 秘密份额上进行计算。
GPU 将计算结果（INT8 秘密份额）返回给参与方。
参与方重构结果。

因此，我们可以预期，本方案的性能将比传统方案有数量级的提升。

4.4 与其他方案的对比

方案	优点	缺点
本方案	1. 深度融合 MPC 与 GPU 计算，消除了交互开销。 2. 充分利用 DeepGEMM 的优化和 GPU 架构特性。 3. 采用 INT8/FP8 低精度计算。 4. GPU 并行 Beaver 三元组生成。 5. JIT 编译优化，kernel 代码高度定制化。 6. 安全性基于信息论安全的秘密共享。	1. 需要对 DeepGEMM kernel 进行深度重构，开发难度较高。 2. 安全性依赖于 GPU 不泄露秘密份额（半诚实模型）。 3. 目前主要支持加法秘密共享和 Beaver 三元组乘法，对其他 MPC 协议的支持需要进一步研究。
传统“两张皮”MPC-GEMM 方案	1. MPC 协议与 GPU 计算分离，模块化程度高，易于实现和维护。 2. 可以使用现有的 MPC 框架和 GPU 加速库。	1. 存在 MPC 协议与 GPU 计算之间的交互开销（数据转换、通信）。 2. GPU 计算部分受到 MPC 协议的制约，无法充分发挥 GPU 的性能和 DeepGEMM 的优化。
基于 TEE 的 MPC-GEMM 方案	1. TEE 提供了一个可信的执行环境，可以保护计算过程的安全性。 2. 可以利用 TEE 内部的 GPU 进行加速计算。	1. 安全性依赖于 TEE 硬件的安全性假设（存在侧信道攻击等风险）。 2. TEE 的性能通常低于原生 GPU 计算。 3. TEE 的可用资源（内存、计算能力）有限。 4. 不同厂商的 TEE 实现存在差异，可移植性较差。
基于同态加密的 MPC-GEMM 方案	1. 安全性基于数学难题（如格密码），安全性高。 2. 可以在密文上直接进行计算，无需解密。	1. 计算开销非常大，通常比明文计算慢几个数量级，难以应用于大规模矩阵运算。 2. 通信开销也很大，因为密文通常比明文大很多。 3. 支持的运算类型有限，通常只支持加法和乘法同态，难以支持复杂的非线性运算。 4. 需要针对同态加密的特性对算法和 kernel 进行重新设计。

对比总结：

性能：本方案 > 基于 TEE 的方案 > 传统“两张皮”方案 > 基于同态加密的方案
安全性：本方案 ≈ 基于同态加密的方案 > 传统“两张皮”方案 > 基于 TEE 的方案
开发难度：基于同态加密的方案 > 本方案 > 基于 TEE 的方案 > 传统“两张皮”方案
硬件依赖：基于 TEE 的方案 > 本方案 > 传统“两张皮”方案 ≈ 基于同态加密的方案
灵活性：传统“两张皮”方案 > 本方案 > 基于 TEE 的方案 > 基于同态加密的方案

5. 总结

本文提出了一种基于秘密共享重构 DeepGEMM kernel 的 MPC-GEMM 方案。该方案通过将 MPC 协议逻辑直接嵌入到 DeepGEMM kernel 中，实现了 MPC 与 GPU 计算的深度融合，彻底消除了传统方案中的“两张皮”问题。方案充分利用了 DeepGEMM 的优化技术、Hopper 架构特性、INT8/FP8 低精度计算、GPU 并行 Beaver 三元组生成以及 JIT 编译等关键技术，在保证计算安全性的前提下，最大程度地发挥了 GPU 的计算能力。

相比于传统的 MPC-GEMM 方案，理论上本方案在性能上具有显著优势，同时在安全性方面也达到了较高的水平。本方案为构建高效安全的 MPC-GEMM 提供了一条全新的技术路线，是对 MPC 与 GPU 加速深度融合的一次探索性设想。

参考链接：https://github.com/deepseek-ai/DeepGEMM

本文转载自上堵吟，作者：上堵吟

标签

DeepSeek

Kernel

MPC

已于2025-3-13 16:17:06修改

51CTO

51CTO博客

51CTO学堂