摘要

本文針對安全多方計算（MPC）框架下通用矩陣乘法（GEMM）運算的性能瓶頸，提出一種全新的 MPC-GEMM 實現(xiàn)方案。該方案的核心思想在于：基于加法秘密共享重構(gòu) DeepSeek DeepGEMM 的 CUDA kernel，將 MPC 協(xié)議的邏輯與 DeepGEMM 的底層優(yōu)化深度融合，消除 MPC 協(xié)議與 GPU 計算之間的“兩張皮”現(xiàn)象。方案采用 INT8/FP8 數(shù)據(jù)表示、秘密共享運算的 kernel 級實現(xiàn)、Hopper 架構(gòu)優(yōu)化（如適用）、GPU 并行 Beaver 三元組生成以及 JIT 編譯等關(guān)鍵技術(shù)。本文將詳細闡述方案的設(shè)計原理、技術(shù)框架、實現(xiàn)細節(jié)（包括 kernel 代碼示例、算法描述、優(yōu)化策略），并從可行性、安全性、高效性等方面進行全面深入的論證，最后與其他 MPC-GEMM 方案進行對比。方案旨在實現(xiàn)真正意義上的安全、高效的 MPC-GEMM，為隱私保護機器學(xué)習(xí)提供強有力的支持。

關(guān)鍵詞： DeepGEMM, DeepSeek, MPC, GEMM, 秘密共享, CUDA, Kernel 重構(gòu), 安全計算, INT8, FP8, Hopper 架構(gòu), Beaver 三元組, JIT 編譯, 并行計算

1. 引言：MPC-GEMM 的性能挑戰(zhàn)與 DeepGEMM 的機遇

安全多方計算（MPC）使得互不信任的參與方能夠在不泄露各自私有數(shù)據(jù)的前提下進行協(xié)同計算，是實現(xiàn)隱私保護機器學(xué)習(xí)的關(guān)鍵技術(shù)。通用矩陣乘法（GEMM）作為深度學(xué)習(xí)模型的核心運算，其在 MPC 框架下的實現(xiàn)（MPC-GEMM）的效率直接影響著隱私保護機器學(xué)習(xí)應(yīng)用的整體性能和實用性。然而，現(xiàn)有的 MPC-GEMM 方案普遍面臨著嚴重的性能挑戰(zhàn)：

• 計算開銷：MPC 協(xié)議的密碼學(xué)運算（如秘密共享、同態(tài)加密）計算復(fù)雜度遠高于明文計算。
• 通信開銷：多數(shù) MPC 協(xié)議需要在參與方之間進行大量的交互通信，尤其是在執(zhí)行乘法運算時，通信開銷成為主要瓶頸。
• 硬件加速：如何在 MPC 的安全約束下有效利用 GPU 等硬件加速器進行計算，是一個極具挑戰(zhàn)性的問題。

傳統(tǒng)的 MPC-GEMM 方案通常采用“兩張皮”模式：MPC 協(xié)議負責保證計算的安全性，GPU 負責提供計算加速，兩者之間通過某種安全接口（如可信執(zhí)行環(huán)境 TEE 或同態(tài)加密）進行交互。這種模式的缺點在于：

• 交互開銷：MPC 協(xié)議與 GPU 計算之間存在數(shù)據(jù)轉(zhuǎn)換（如明文與密文、秘密份額與 GPU 可處理格式之間的轉(zhuǎn)換）和通信的開銷，限制了整體性能。
• GPU 利用率：GPU 計算部分通常受到 MPC 協(xié)議的制約，無法充分發(fā)揮 GPU 的并行計算能力和 DeepGEMM 等底層優(yōu)化庫的性能優(yōu)勢。

DeepSeek 最新發(fā)布的 DeepGEMM 是一個為 NVIDIA GPU 優(yōu)化的高性能 GEMM 庫。它通過 FP8 低精度計算、針對 GPU 架構(gòu)的優(yōu)化、CUDA kernel 優(yōu)化以及 JIT 編譯等技術(shù)，大幅提升了 GEMM 運算的效率。雖然 DeepGEMM 并非專門為 MPC 設(shè)計，但其在 kernel 級別的優(yōu)化為我們提供了一個重要的機遇：能否將 MPC 協(xié)議與 DeepGEMM 的底層優(yōu)化進行深度融合，消除“兩張皮”現(xiàn)象，實現(xiàn)真正意義上的安全高效的 MPC-GEMM？

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

基于 MPC 與 DeepGEMM 的深度融合，就可以嘗試構(gòu)想一種全新的 MPC-GEMM 方案：基于秘密共享重構(gòu) DeepSeek DeepGEMM kernel。該方案的核心思想是：將 MPC 協(xié)議中與 GEMM 運算相關(guān)的計算邏輯（秘密份額的加法、乘法）直接實現(xiàn)在 DeepGEMM 的 CUDA kernel 中，讓 GPU 直接執(zhí)行一個完整的“MPC-GEMM”運算。

方案的設(shè)計基于以下幾個關(guān)鍵原理：

1）加法秘密共享：采用加法秘密共享作為 MPC 的基礎(chǔ)安全機制。加法秘密共享具有以下優(yōu)點：

• 簡單高效：實現(xiàn)簡單，只需要進行模加運算。
• 加法同態(tài)：秘密份額的加法對應(yīng)于明文的加法，使得加法運算可以在本地高效執(zhí)行，無需通信。
• 安全性：信息論安全，只要參與方不合謀，任何單獨的秘密份額都不會泄露關(guān)于原始數(shù)據(jù)的任何信息。

2）INT8/FP8 數(shù)據(jù)表示：為了降低計算和通信開銷，我們借鑒 DeepGEMM 對低精度計算的使用，將輸入數(shù)據(jù)（FP32/FP64/定點數(shù)）映射到 INT8 或 FP8。

• INT8 映射：對于 INT8，我們采用偏移映射等策略，充分利用 INT8 的表示范圍，并簡化秘密共享運算。
• FP8 映射:如果采用 FP8，可以利用 DeepGEMM 自身的 FP8 支持。

3）DeepGEMM Kernel 重構(gòu)：方案的核心在于對 DeepGEMM 的 CUDA kernel 進行重構(gòu)。我們將 MPC 協(xié)議的邏輯（即秘密共享下的加法和乘法）直接嵌入到 kernel 中。

• 輸入/輸出：Kernel 的輸入和輸出直接是秘密份額（INT8 或 FP8），而不是明文數(shù)據(jù)。
• 基本運算：將 kernel 中的加法和乘法替換為 MPC 協(xié)議下的秘密共享加法和乘法（基于 Beaver 三元組）。
• 保留優(yōu)化：盡最大可能保留 DeepGEMM 原有的針對 GPU 架構(gòu)的優(yōu)化技術(shù)，如 tiling、loop unrolling、shared memory 利用、warp-level primitives、指令級并行等，并針對秘密共享運算進行適配。
• 異步計算: 盡可能利用GPU的異步計算能力。

4）Beaver 三元組乘法：為了在秘密共享下實現(xiàn)乘法，采用 Beaver 三元組乘法協(xié)議?？梢栽?kernel 中實現(xiàn) Beaver 三元組乘法協(xié)議，并利用 warp-level primitives（如??__shfl_xor_sync??）進行優(yōu)化。

5）GPU 并行 Beaver 三元組生成：為了提高 Beaver 三元組的生成效率，并減少預(yù)處理階段的通信開銷，我們可以利用 GPU 的并行計算能力，在 GPU 上并行生成 Beaver 三元組。

6）JIT 編譯：我們充分利用 DeepGEMM 的 JIT 編譯技術(shù)（如果 DeepGEMM 提供 JIT 編譯接口；如果沒有，我們可以自行實現(xiàn) JIT 編譯），根據(jù) GEMM 形狀、塊大小、參與方數(shù)量等參數(shù)，動態(tài)生成高度優(yōu)化的 MPC-GEMM kernel。

7）簡化的 MPC 協(xié)議：由于 GPU 直接參與 MPC 協(xié)議的執(zhí)行（我們將其視為一個“半誠實”的參與方），我們可以簡化 MPC 協(xié)議的設(shè)計，減少通信輪數(shù)和通信量。

3. 技術(shù)框架與實現(xiàn)細節(jié)

3.1 技術(shù)框架

方案的技術(shù)框架主要由以下幾個模塊構(gòu)成：

• 秘密共享模塊：

a.負責將參與方的輸入數(shù)據(jù)（FP32、FP64 或定點數(shù)）進行加法秘密共享。

b.將秘密份額轉(zhuǎn)換為 INT8 或 FP8 表示（通過映射）。

c.實現(xiàn)秘密共享上的加法和乘法運算（基于 Beaver 三元組）。

d.提供秘密份額的生成、分發(fā)、重構(gòu)等功能。

• DeepGEMM Kernel 重構(gòu)模塊：

a.負責對 DeepGEMM 的 CUDA kernel 進行重構(gòu)，將秘密共享運算（加法和乘法）嵌入到 kernel 中。

b.保留并適配 DeepGEMM 原有的 GPU 架構(gòu)優(yōu)化。

c.利用 JIT 編譯技術(shù)（或手動實現(xiàn)），動態(tài)生成針對特定參數(shù)（GEMM 形狀、塊大小、參與方數(shù)量等）的優(yōu)化 kernel。

• MPC 協(xié)議協(xié)調(diào)模塊：

a.負責協(xié)調(diào)各參與方和 GPU 之間的交互。

b.管理 Beaver 三元組的分發(fā)（如果采用離線生成）。

c.觸發(fā) GPU kernel 的執(zhí)行。

• GPU Beaver 三元組生成模塊：

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

3.2 工作流程

整個 MPC-GEMM 的計算流程分為離線階段和在線階段：

1. 離線階段（預(yù)處理）：

• 利用 GPU 并行生成 Beaver 三元組，并將三元組的秘密份額分發(fā)給各參與方（和 GPU 線程）。

2. 在線階段：

• 參與方收集各自的輸出份額。
• 將對應(yīng)位置的份額相加（模運算，如果是 INT8；浮點加法，如果是 FP8），重構(gòu)出最終的 GEMM 結(jié)果。
• 如果需要，可以將結(jié)果轉(zhuǎn)換回 FP32 或 FP64 格式。
• Kernel 計算完成后，輸出結(jié)果仍然是秘密份額（INT8 或 FP8）的形式。
• GPU 將輸出份額返回給參與方。
• GPU 執(zhí)行重構(gòu)后的 DeepGEMM kernel。
• 在 kernel 內(nèi)部：
• 整個計算過程高度并行化。
• 將輸入數(shù)據(jù)（秘密份額）和 Beaver 三元組份額加載到 shared memory。
• 使用 tiling 技術(shù)將矩陣分塊。
• 對于每個塊，執(zhí)行秘密共享下的加法和乘法運算（利用 Beaver 三元組和 warp-level primitives）。
• 利用 GPU 架構(gòu)優(yōu)化（如 tiling, loop unrolling, shared memory, warp-level primitives, 指令級并行, 異步計算等）。
• 將中間結(jié)果累加到 shared memory 或 registers 中。
• MPC 協(xié)議協(xié)調(diào)模塊根據(jù) GEMM 運算的參數(shù)（形狀、塊大小等）和參與方數(shù)量，觸發(fā) DeepGEMM Kernel 重構(gòu)模塊生成相應(yīng)的 CUDA kernel（利用 JIT 編譯或手動實現(xiàn)）。
• 參與方將各自持有的秘密份額（INT8 或 FP8）直接作為輸入，傳遞給生成的 CUDA kernel。
• 每個參與方將自己的輸入矩陣的每個元素進行加法秘密共享。
• 將秘密份額轉(zhuǎn)換為 INT8 或 FP8 表示。
• ?輸入準備：
• Kernel 調(diào)用：
• GPU 并行計算：
• 輸出處理：
• 結(jié)果重構(gòu)：?

3.3 關(guān)鍵實現(xiàn)細節(jié)

本文中所用代碼均是偽代碼，根據(jù)通義靈碼的建議生成的，只能看出大致的意思，不能直接使用。

3.3.1 數(shù)據(jù)表示

• INT8 映射 (如果采用 INT8)：
  我們推薦使用偏移映射。假設(shè)原始數(shù)據(jù)為 FP32，映射規(guī)則如下：
  映射公式：
  其中，S1、S2 是縮放因子，O1、O2 是偏移量。具體數(shù)值需要根據(jù)實際數(shù)據(jù)分布和 INT8 的表示范圍來確定。
• 對于 FP32 正數(shù) x：??INT8 = round(x * S1) + O1??
• 對于 FP32 負數(shù) x：??INT8 = round(x * S2) + O2??
• 將 FP32 的 NaN 映射到 INT8 的 -128。
• 將 FP32 的 +Inf 映射到 INT8 的 -127。
• 將 FP32 的 -Inf 映射到 INT8 的 -126。
• 將 FP32 的 0 映射到 INT8 的 0。
• 將 FP32 的其他正數(shù)，等比例映射到 INT8 的 [1, 127] 區(qū)間。
• 將 FP32 的其他負數(shù)，等比例映射到 INT8 的 [-125, -1] 區(qū)間。
• FP8 表示 (如果采用 FP8):如果采用FP8，可以直接利用DeepGEMM對FP8的支持。

3.3.2 CUDA Kernel 中的秘密共享乘法

以下是 CUDA kernel 中實現(xiàn)秘密共享乘法（基于加法秘密共享和 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;

// 循環(huán)處理矩陣塊 (tiling)
    int buffer_idx = 0; // 雙緩沖索引
    for (int i = 0; i < k; i += BLOCK_SIZE) {
      // 將數(shù)據(jù)從全局內(nèi)存加載到 shared memory (異步加載, 如果支持)
       if (grid.rank() == 0 && i + BLOCK_SIZE < k) {
            //僅rank 0 的block進行異步加載
            //這里只是偽代碼，實際使用需要根據(jù)數(shù)據(jù)類型進行調(diào)整
            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(); // 等待所有線程加載完成, 以及異步加載完成

      // 計算當前塊的乘積 (循環(huán)展開)
      #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;
          
        //根據(jù)數(shù)據(jù)類型進行模運算
        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;
    }
    // 將結(jié)果寫回全局內(nèi)存
    if(row < m && col < n){
        z_shares[row * n + col] = acc;
    }
}

代碼解釋：

• 模板參數(shù) ??T??：使用模板參數(shù) ??T??，可以支持 INT8 和 FP8 兩種數(shù)據(jù)類型。
• 雙緩沖 (Double Buffering)：使用兩組 shared memory 數(shù)組，實現(xiàn)計算和數(shù)據(jù)加載的流水線操作。
• 異步數(shù)據(jù)加載：在外層循環(huán)的開始處，嘗試使用 ??cudaMemcpyAsync?? 異步地將下一批次的數(shù)據(jù)從全局內(nèi)存加載到 shared memory。
• Cooperative Groups：使用 Cooperative Groups 提供的 ??thread_block???、??grid_group??? 和 ??thread_block_tile?? 類型來更精細地控制線程塊和 warp 級別的并行。
• Warp-level Shuffle 指令優(yōu)化：

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

    b.循環(huán)展開 warp-level shuffle 操作。

• 指令級并行（手動）：在計算 ??z?? 時，將乘法和加法運算交錯進行，盡可能利用 GPU 的指令級并行能力。
• 循環(huán)展開：使用 ??#pragma unroll?? 指令展開內(nèi)層循環(huán)。
• 模運算: 如果 ??T??? 是 ??int8_t???，則使用 ??& 0xFF?? 進行模 256 運算。
• Tiling: 使用tiling技術(shù)將矩陣分塊處理。
• 并行化:

 a.線程塊 (Block)：不同的線程塊負責計算輸出矩陣 Z 的不同塊（tiling）。

 b.線程 (Thread)：線程塊內(nèi)部的線程協(xié)同計算秘密共享乘法。

3.3.3 Hopper 架構(gòu)優(yōu)化（深化）

• TMA (Tensor Memory Accelerator)：通過流水線、雙緩沖和 ??cudaMemcpyAsync??，盡可能利用 TMA 的異步數(shù)據(jù)傳輸能力，隱藏內(nèi)存訪問延遲。
• 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[...], ...); // 加載數(shù)據(jù)到 fragment, 需要根據(jù)實際情況填寫參數(shù)
     wmma::load_matrix_sync(b_frag, &y_shared[...], ...); // 加載數(shù)據(jù)到 fragment, 需要根據(jù)實際情況填寫參數(shù)

     wmma::mma_sync(acc_frag, a_frag, b_frag, acc_frag); // 矩陣乘累加
 }
//將秘密共享乘法結(jié)果加到acc_frag上
 wmma::store_matrix_sync(&c_shared[...], acc_frag, ... , wmma::mem_row_major); // 存儲結(jié)果
 ```
*   **FP8 計算:** 如果采用FP8, 可以直接使用DeepGEMM中針對FP8和Tensor Core的優(yōu)化。
*   **數(shù)據(jù)類型轉(zhuǎn)換：** 如果 INT8 的 `wmma.mma` 指令效果不佳，可考慮將 INT8 份額轉(zhuǎn)換為 FP16 或 INT32，然后使用相應(yīng)的 `wmma.mma` 指令。但類型轉(zhuǎn)換也需在秘密共享下進行。

 a.INT8 計算：嘗試使用 ??wmma::mma_s8s8s32?? 指令進行 INT8 矩陣乘法：

• Shared Memory 優(yōu)化：

a.通過 tiling 技術(shù)和合理的數(shù)據(jù)訪問模式，最大程度地復(fù)用 shared memory 中的數(shù)據(jù)。

b.合理安排 shared memory 中數(shù)據(jù)的存儲位置，避免 bank conflict。

• Warp-level Primitives 與指令級并行：

a.充分利用??__shfl_xor_sync??? 或??warp.shfl_xor?? 指令在 warp 內(nèi)部高效地進行數(shù)據(jù)交換和規(guī)約求和。

b.在 kernel 代碼中，盡可能地將獨立的指令放在一起執(zhí)行，利用 GPU 的指令級并行能力。

3.3.4 GPU 并行 Beaver 三元組生成

算法：

1. 初始化 cuRAND：在每個線程中初始化一個 cuRAND 偽隨機數(shù)生成器狀態(tài)。
2. 生成隨機數(shù)：使用 cuRAND 庫在每個線程中并行生成三個 INT8 或 FP8 類型的隨機數(shù)（a, b, c）。
3. 驗證三元組：在每個線程中驗證生成的三元組是否滿足 Beaver 三元組的條件（??c == a * b??）。
4. 秘密共享：在 kernel 中直接對驗證通過的三元組 (a, b, c) 進行加法秘密共享。
5. 存儲份額：將每個參與方的三元組份額存儲到全局內(nèi)存中的一個數(shù)組中。

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 偽隨機數(shù)生成器
    curandState_t state;
    curand_init(blockIdx.x * blockDim.x + threadIdx.x, 0, 0, &state);

    // 生成 Beaver 三元組并進行秘密共享
    if (tid < num_triples) {
        // 1. 生成隨機數(shù)
        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 {
           // 如果驗證失敗，可以將其設(shè)置為一個特殊值（如全 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 庫的函數(shù)聲明。
• ??BeaverTripleShares?? 結(jié)構(gòu)體：定義了 Beaver 三元組份額的結(jié)構(gòu)。
• ??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???：指向全局內(nèi)存中存儲 Beaver 三元組份額的數(shù)組的指針，其大小應(yīng)為 ??num_triples * num_parties??。
• ??num_triples??：要生成的 Beaver 三元組的數(shù)量。
• ??num_parties??：參與方的數(shù)量。
• ??tid??：線程 ID。
• ??curandState_t??：cuRAND 偽隨機數(shù)生成器的狀態(tài)。每個線程都需要一個獨立的狀態(tài)。
• ??curand_init??：初始化偽隨機數(shù)生成器。這里使用線程 ID 作為種子，確保每個線程生成的隨機數(shù)序列不同。
• ??curand???：生成一個 32 位無符號整數(shù)隨機數(shù)。我們將其強制轉(zhuǎn)換為 ??int8??。
• 驗證三元組：
• 秘密共享：
• 存儲份額：將每個參與方的三元組份額存儲到 ??triples?? 數(shù)組中。

使用方法：

• 在 GPU 上分配足夠大的內(nèi)存來存儲 Beaver 三元組的所有份額 (??BeaverTripleShares* triples??)。
• 調(diào)用 ??generate_beaver_triples?? kernel，生成 Beaver 三元組并進行秘密共享。
• 在 MPC-GEMM kernel 中，每個線程根據(jù)其線程 ID 和參與方 ID 從 ??triples?? 數(shù)組中獲取相應(yīng)的 Beaver 三元組份額。

優(yōu)化：

• 可以通過增加線程塊和線程數(shù)量來進一步提高 Beaver 三元組生成的并行度。
• 可以使用更高效的隨機數(shù)生成器（如 Philox 算法）來提高隨機數(shù)生成的速度和質(zhì)量。
• 可以將 Beaver 三元組的生成、驗證和秘密共享融合到一個 kernel 中，減少數(shù)據(jù)傳輸開銷。

3.3.5 JIT 編譯優(yōu)化

JIT 編譯技術(shù)允許我們在運行時根據(jù)具體的參數(shù)動態(tài)生成優(yōu)化的 CUDA kernel 代碼。在 MPC-GEMM 中，我們可以利用 JIT 編譯進行以下優(yōu)化：

1. 代碼特化：

• GEMM 參數(shù)：根據(jù) GEMM 運算的形狀（M, N, K）、塊大小（BLOCK_SIZE）、數(shù)據(jù)類型（INT8 或 FP8）等參數(shù)，生成專門針對這些參數(shù)優(yōu)化的 kernel 代碼。例如，可以根據(jù) M、N、K 的大小選擇最合適的 tiling 策略和 shared memory 使用方式。
• MPC 參數(shù)：根據(jù)參與方數(shù)量、秘密共享方案（加法秘密共享）等參數(shù)，生成相應(yīng)的 kernel 代碼。例如，如果參與方數(shù)量較少，可以使用更激進的 warp-level shuffle 優(yōu)化。
• Hopper 架構(gòu)特性：根據(jù)目標 GPU 的計算能力（Compute Capability），啟用或禁用某些 Hopper 架構(gòu)特有的優(yōu)化（如 TMA）。

2. 常量折疊：

• Beaver 三元組內(nèi)聯(lián)：如果 Beaver 三元組是在預(yù)處理階段生成的，并且在 kernel 執(zhí)行期間不會改變，可以將三元組的份額直接作為編譯時常量內(nèi)聯(lián)到 kernel 代碼中，減少運行時內(nèi)存訪問。
• 其他常量：將參與方數(shù)量、塊大小、GEMM 形狀等參數(shù)也作為編譯時常量內(nèi)聯(lián)到 kernel 代碼中，允許編譯器進行更多的優(yōu)化（如常量傳播、死代碼消除等）。

3. 循環(huán)展開：

• 根據(jù) GEMM 形狀和塊大小，對 kernel 中的循環(huán)進行部分或完全展開，減少循環(huán)控制開銷，并增加指令級并行度。
• 特別是對于秘密共享乘法協(xié)議中的內(nèi)層循環(huán)，可以進行更激進的展開。

4. 指令級并行：

• JIT 編譯器可以分析 kernel 代碼中的數(shù)據(jù)依賴關(guān)系，盡可能地將獨立的指令放在一起執(zhí)行，利用 GPU 的指令級并行能力。
• 我們可以手動調(diào)整 kernel 代碼中的指令順序，以幫助編譯器更好地進行指令級并行優(yōu)化。

5. 自動調(diào)整block size和grid size: 可以根據(jù)矩陣規(guī)模、數(shù)據(jù)類型等，自動調(diào)整kernel的block size和grid size，以充分利用GPU資源。

實現(xiàn)方式：

• NVRTC (NVIDIA Runtime Compilation)：NVRTC 是 NVIDIA 提供的一個運行時編譯庫，可以在程序運行時將 CUDA C++ 代碼編譯為 PTX 匯編代碼，然后加載到 GPU 中執(zhí)行。
• NVCC (NVIDIA CUDA Compiler)：NVCC 是 NVIDIA 的 CUDA 編譯器，也可以用于 JIT 編譯?？梢栽诰幾g時使用 ??-D??? 選項定義宏，然后在 kernel 代碼中使用 ??#ifdef?? 等預(yù)處理指令來根據(jù)不同的宏定義生成不同的代碼。

示例：

假設(shè)我們要根據(jù)參與方數(shù)量 ??n?? 進行代碼特化。我們可以在 kernel 代碼中使用如下預(yù)處理指令：

#if N_PARTIES == 2
    // 針對 2 個參與方的優(yōu)化代碼
    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 個參與方的優(yōu)化代碼
     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 就會生成針對特定參與方數(shù)量的優(yōu)化 kernel 代碼。

4. 方案論證

4.1 可行性論證

• DeepGEMM Kernel 可修改性：DeepGEMM 的 CUDA kernel 本質(zhì)上是 C/C++ 代碼，可以進行修改和擴展。
• 秘密共享運算可實現(xiàn)性：加法秘密共享和基于 Beaver 三元組的乘法協(xié)議都可以在 INT8 或 FP8 數(shù)據(jù)類型上高效實現(xiàn)。
• GPU 并行計算可行性：CUDA 編程模型支持細粒度的并行計算，可以充分利用 GPU 的并行計算能力。
• JIT 編譯可行性：JIT 編譯技術(shù)已經(jīng)廣泛應(yīng)用，NVRTC 和 NVCC 都提供了 JIT 編譯功能。
• Hopper 架構(gòu)優(yōu)化可行性：Hopper 架構(gòu)的特性（TMA、Tensor Core、Shared Memory、Warp-level Primitives）都可以在 CUDA 編程中加以利用。

4.2 安全性論證

本方案的安全性基于以下幾個方面：

1. 秘密共享的安全性：采用的加法秘密共享方案是信息論安全的，只要參與方不合謀，任何單獨的秘密份額都不會泄露關(guān)于原始數(shù)據(jù)的任何信息。
2. Beaver 三元組乘法協(xié)議的安全性：Beaver 三元組乘法協(xié)議在半誠實模型下是安全的。只要 Beaver 三元組是獨立于輸入數(shù)據(jù)生成的，并且參與方誠實地執(zhí)行協(xié)議，攻擊者就無法從公開的中間值（d 和 e）中推斷出關(guān)于秘密輸入（x 和 y）的任何信息。
3. GPU 計算的安全性：

• GPU 始終只接觸到秘密份額，無法獲得任何關(guān)于明文數(shù)據(jù)的信息。
• 重構(gòu)后的 DeepGEMM kernel 只執(zhí)行秘密共享運算，不包含任何可能泄露敏感信息的操作（如直接訪問內(nèi)存地址、向外部發(fā)送數(shù)據(jù)等）。
• 即使攻擊者控制了 GPU，也只能獲得秘密份額，無法恢復(fù)出原始數(shù)據(jù)。

4. JIT 編譯的安全性：

• JIT 編譯器生成的 kernel 代碼只包含必要的秘密共享運算和優(yōu)化邏輯，不包含任何惡意代碼。
• 可以對 JIT 編譯器生成的代碼進行靜態(tài)分析和安全審計。

5. 抵御側(cè)信道攻擊：

• 雖然 GPU 內(nèi)部的計算對參與方透明，但仍然需要考慮側(cè)信道攻擊（如時間攻擊、功耗攻擊）。
• 可以采用掩碼（masking）技術(shù)來防御側(cè)信道攻擊。具體來說，可以將秘密份額與一個隨機數(shù)進行運算（如異或），然后在掩碼后的份額上進行計算，最后再去除掩碼。
• 可以對 kernel 代碼進行隨機化，使得每次執(zhí)行的指令順序和內(nèi)存訪問模式都不同，增加側(cè)信道攻擊的難度。

4.3 高效性論證

相比于傳統(tǒng)的“兩張皮”MPC-GEMM 方案，本方案具有以下優(yōu)勢：

• 消除交互開銷：將 MPC 協(xié)議邏輯直接嵌入到 DeepGEMM kernel 中，徹底消除了 MPC 協(xié)議與 GPU 計算之間的所有交互開銷（如數(shù)據(jù)格式轉(zhuǎn)換、安全通道傳輸?shù)龋?。這是本方案相對于傳統(tǒng)方案最大的優(yōu)勢所在。
• 充分利用 DeepGEMM 優(yōu)化：GPU 直接執(zhí)行 MPC-GEMM 運算，可以充分利用 DeepGEMM 原有的針對 GPU 架構(gòu)的各種優(yōu)化（tiling、loop unrolling、shared memory 利用、TMA、Tensor Core、warp-level primitives、指令級并行等）。
• 低精度計算：使用 INT8 或 FP8 數(shù)據(jù)類型，相比于 FP32 或 FP64，可以顯著減少計算量和通信量。
• GPU 并行 Beaver 三元組生成：利用 GPU 并行生成 Beaver 三元組，大幅減少了預(yù)處理階段的開銷。
• 簡化的 MPC 協(xié)議：將 GPU 視為“半誠實”參與方，可以簡化 MPC 協(xié)議的設(shè)計，減少通信輪數(shù)。
• JIT 編譯優(yōu)化：通過 JIT 編譯，可以針對具體的 GEMM 參數(shù)和 MPC 參數(shù)生成高度定制化的 kernel 代碼，進一步提升性能。
• 高度并行化: 秘密共享的加法、乘法，Beaver三元組的生成都可以在GPU上高度并行。

量化分析（舉例）：

假設(shè)一個 MPC-GEMM 運算涉及兩個矩陣 A 和 B 的乘法，矩陣大小為 1024x1024，參與方數(shù)量為 3。

傳統(tǒng)“兩張皮”方案：
整個過程中，數(shù)據(jù)至少需要在網(wǎng)絡(luò)上傳輸 3 次（輸入 2 次，輸出 1 次），并且涉及到多次數(shù)據(jù)格式轉(zhuǎn)換。

• 參與方之間需要通過網(wǎng)絡(luò)傳輸秘密份額（FP32 或 FP64）。
• 需要將秘密份額轉(zhuǎn)換為 GPU 可處理的格式（如加密）。
• GPU 執(zhí)行 GEMM 計算。
• 將計算結(jié)果（加密或編碼）傳輸回參與方。
• 參與方進行解密或解碼，并重構(gòu)結(jié)果。

本方案：
整個過程中，數(shù)據(jù)只需要在網(wǎng)絡(luò)上傳輸 2 次（輸入和輸出），并且都是 INT8 類型，數(shù)據(jù)量大大減少。GPU 內(nèi)部的計算高度優(yōu)化，且無需與 MPC 協(xié)議進行交互。

• 參與方將輸入數(shù)據(jù)進行秘密共享，并映射到 INT8。
• 參與方將 INT8 秘密份額直接發(fā)送給 GPU（通過 MPI 等）。
• GPU 執(zhí)行重構(gòu)后的 DeepGEMM kernel，直接在 INT8 秘密份額上進行計算。
• GPU 將計算結(jié)果（INT8 秘密份額）返回給參與方。
• 參與方重構(gòu)結(jié)果。

因此，我們可以預(yù)期，本方案的性能將比傳統(tǒng)方案有數(shù)量級的提升。

4.4 與其他方案的對比

對比總結(jié)：

• 性能：本方案 > 基于 TEE 的方案 > 傳統(tǒng)“兩張皮”方案 > 基于同態(tài)加密的方案
• 安全性：本方案 ≈ 基于同態(tài)加密的方案 > 傳統(tǒng)“兩張皮”方案 > 基于 TEE 的方案
• 開發(fā)難度：基于同態(tài)加密的方案 > 本方案 > 基于 TEE 的方案 > 傳統(tǒng)“兩張皮”方案
• 硬件依賴：基于 TEE 的方案 > 本方案 > 傳統(tǒng)“兩張皮”方案 ≈ 基于同態(tài)加密的方案
• 靈活性：傳統(tǒng)“兩張皮”方案 > 本方案 > 基于 TEE 的方案 > 基于同態(tài)加密的方案

5. 總結(jié)

本文提出了一種基于秘密共享重構(gòu) DeepGEMM kernel 的 MPC-GEMM 方案。該方案通過將 MPC 協(xié)議邏輯直接嵌入到 DeepGEMM kernel 中，實現(xiàn)了 MPC 與 GPU 計算的深度融合，徹底消除了傳統(tǒng)方案中的“兩張皮”問題。方案充分利用了 DeepGEMM 的優(yōu)化技術(shù)、Hopper 架構(gòu)特性、INT8/FP8 低精度計算、GPU 并行 Beaver 三元組生成以及 JIT 編譯等關(guān)鍵技術(shù)，在保證計算安全性的前提下，最大程度地發(fā)揮了 GPU 的計算能力。

相比于傳統(tǒng)的 MPC-GEMM 方案，理論上本方案在性能上具有顯著優(yōu)勢，同時在安全性方面也達到了較高的水平。本方案為構(gòu)建高效安全的 MPC-GEMM 提供了一條全新的技術(shù)路線，是對 MPC 與 GPU 加速深度融合的一次探索性設(shè)想。

參考鏈接：??https://github.com/deepseek-ai/DeepGEMM??

本文轉(zhuǎn)載自??上堵吟??，作者：上堵吟