Scaling up Privacy-Preserving ML: A CKKS Implementation of Llama-2-7B

TL;DR

This work addresses the privacy risks of LLM inference by proposing CKKS-based private inference for Llama models with long input contexts. It introduces a heterogeneous unbalanced chunked prefill framework, along with new homomorphic matrix multiplication and polynomial evaluation techniques, to manage outliers and the blow-up from ciphertext-ciphertext computations. By combining ML-driven range/precision tuning, activation prefixing, and orthogonal rotations, the authors enable end-to-end private inference for up to 4096 input tokens (last 128 encrypted) with practical runtimes on a GPU cluster. The approach delivers a near-bootstrapping-free attention path and a slim Softmax strategy, achieving 85 seconds for summarization and 33 seconds per output token on 8 RTX-4090 GPUs, demonstrating a path toward scalable privacy-preserving LLM inference in real-world deployments.

Abstract

As large language models (LLMs) become ubiquitous, privacy concerns pertaining to inference inputs keep growing. In this context, fully homomorphic encryption (FHE) has emerged as a primary cryptographic solution to provide non-interactive confidential LLM inference. Existing solutions scale poorly with the input token length, and hence focus either on small models or larger models with a small number of input tokens. They also suffer from the existence of large outlier values. These values have a strong impact on the evaluation of non-linear layers, leading to large-degree polynomial approximation and thus heavy evaluation costs. We propose an FHE-based private LLM inference solution that allows thousands of input tokens with only a part of them being encrypted: this fits with a scenario where the context is benign and only part of the input is sensitive. To do so, we suggest an unbalanced chunked prefill framework that processes the private and public parts of the input tokens differently. Our framework contains plaintext-plaintext, plaintext-ciphertext and ciphertext-ciphertext computational components. We adopt different strategies and ingredients for each component. We also devise new homomorphic algorithms for specific matrix multiplication and polynomial evaluation tasks encountered during LLM inference. Furthermore, without retraining, we tailor the LLM inference algorithm to reduce the ranges of outlier values: we leverage machine learning strategies (token prepending and rotations) to mitigate the impact of the outliers on non-linear layers. Based on these ingredients, we describe a CKKS-based end-to-end implementation of Llama-2-7B private inference for up to 4096 input tokens, of which the last 128 are encrypted. On a cluster of 8~NVIDIA RTX-4090 GPUs, inference takes 85s for summarization and 33s for generation per output token.

Figures (3)

• Figure 1: Llama model architecture
• Figure 2: Visualization of the proposed plaintext algorithm with heterogeneous chunked prefill stages.
• Figure 3: Overview of encrypted computation. The red and orange symbols indicate the bootstrap placements. The bootstrapping during Softmax evaluation (orange) is performed only in the narrow auxiliary track (see Section \ref{['se:shallow_Softmax']}).

Theorems & Definitions (7)

• lemma 1
• proof
• theorem 1
• proof
• lemma 2
• proof
• theorem 2