摘要

本文針對安全多方計算（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 的計算流程分為離線階段和在線階段：

1. 離線階段（預處理）：

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

2. 在線階段：

• 參與方收集各自的輸出份額。
• 將對應位置的份額相加（模運算，如果是 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 三元組生成

算法：

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

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 編譯進行以下優化：

1. 代碼特化：

• 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）。

2. 常量折疊：

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

3. 循環展開：

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

4. 指令級并行：

• JIT 編譯器可以分析 kernel 代碼中的數據依賴關系，盡可能地將獨立的指令放在一起執行，利用 GPU 的指令級并行能力。
• 我們可以手動調整 kernel 代碼中的指令順序，以幫助編譯器更好地進行指令級并行優化。

5. 自動調整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 安全性論證

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

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

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

4. JIT 編譯的安全性：

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

5. 抵御側信道攻擊：

• 雖然 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 與其他方案的對比

對比總結：

• 性能：本方案 > 基于 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??

本文轉載自??上堵吟??，作者：上堵吟