Ramp Up NTT in Record Time using GPU-Accelerated Algorithms and LLM-based Code Generation

TL;DR

This paper tackles the HE/PPML bottleneck by developing a GPU-friendly NTT (GNTT) family that leverages PyTorch and CuPy to accelerate polynomial computations and improve HE performance. It also investigates the ability of large language models to automatically generate GPU-optimized NTT code from CPU-oriented references, benchmarking several models. The key finding is that GNTT2 delivers a substantial ~$62\times$ speedup over CPU implementations, while LLM-generated approaches show potential but vary in quality and generally lag behind the optimized GNTT. The work provides practical guidance for turbocharging PPML with GPU acceleration and offers insights into the current capabilities and limitations of LLM-based cryptographic code generation.

Abstract

Homomorphic encryption (HE) is a core building block in privacy-preserving machine learning (PPML), but HE is also widely known as its efficiency bottleneck. Therefore, many GPU-accelerated cryptographic schemes have been proposed to improve the performance of HE. However, these methods often require complex modifications tailored to specific algorithms and are tightly coupled with specific GPU and operating systems. It is interesting to ask how to generally offer more practical GPU-accelerated cryptographic algorithm implementations. Given the powerful code generation capabilities of large language models (LLMs), we aim to explore their potential to automatically generate practical GPU-friendly algorithm code using CPU-friendly code. In this paper, we focus on number theoretic transform (NTT) -- the core mechanism of HE. We first develop and optimize a GPU-friendly NTT (GNTT) family that exploits PyTorch's fast matrix computation and precomputation, achieving an approximately 62x speedup -- a significant boost over existing ones. Then we explore GPU-friendly code generation using various LLMs, including DeepSeek-R1, OpenAI o1 and o3-mini. We discover many interesting findings throughout the process. For instance, somewhat surprisingly, our experiments demonstrate that DeepSeek-R1 significantly outperforms OpenAI o3-mini and o1, but still cannot beat our optimized protocol. The findings provide valuable insights for turbocharging PPML and enhancing code generation capabilities of LLMs. Codes are available at: https://github.com/LMPC-Lab/GenGPUCrypto.

Figures (6)

• Figure 1: The framework of GPU-Accelerated cryptographic algorithm code generation based on LLMs and its benchmarking. We provide an example of algorithmic code generation on how to convert the core CPU-friendly code for GNTT precomputation into a code that takes into account the computational advantages of both CPU and GPU. This example actually illustrates the core difference between GNTT1 and GNTT2.
• Figure 2: Evaluation results based on parameters $q$ and $n$ that satisfy PQC for all algorithms.
• Figure 3: Evaluation results of the precomputation time for GNTT algorithms.
• Figure 4: Comparison of the evaluation results for NTT running time.
• Figure 5: Comparison of the evaluation results for INTT running time.