AEGIS: Scaling Long-Sequence Homomorphic Encrypted Transformer Inference via Hybrid Parallelism on Multi-GPU Systems

Zhaoting Gong; Ran Ran; Fan Yao; Wujie Wen

cs.CRApr 3, 2026

AEGIS: Scaling Long-Sequence Homomorphic Encrypted Transformer Inference via Hybrid Parallelism on Multi-GPU Systems

Authors: Zhaoting Gong, Ran Ran, Fan Yao, Wujie Wen

Organizations: North Carolina State University, Raleigh, NC, USA · University of Central Florida, Orlando, FL, USA

Abstract

Fully Homomorphic Encryption (FHE) enables privacy-preserving Transformer inference, but long-sequence encrypted Transformers quickly exceed single-GPU memory capacity because encoded weights are already large and encrypted activations grow rapidly with sequence length. Multi-GPU execution therefore becomes unavoidable, yet scaling remains challenging because communication is jointly induced by application-level aggregation and encryption-level RNS coupling. Existing approaches either synchronize between devices frequently or replicate encrypted tensors across devices, leading to excessive communication and latency. We present AEGIS, an Application-Encryption Guided Inference System for scalable long-sequence encrypted Transformer inference on multi-GPU platforms. AEGIS derives device placement from ciphertext dependencies jointly induced by Transformer dataflow and CKKS polynomial coupling, co-locating modulus-coherent and token-coherent data so that communication is introduced only when application dependencies require it, while reordering polynomial operators to overlap the remaining collectives with computation. On 2048-token inputs, AEGIS reduces inter-GPU communication by up to 57.9% in feed-forward networks and 81.3% in self-attention versus prior state-of-the-art designs. On four GPUs, it achieves up to 96.62% scaling efficiency, 3.86x end-to-end speedup, and 69.1% per-device memory reduction. These results establish coordinated application-encryption parallelism as a practical foundation for scalable homomorphic Transformer inference.

Figures & tables

Notation	Description
$N$	Polynomial degree (number of coefficients)
$S$	Number of available slots in an encoded message
$Q_{L}$	Ciphertext modulus chain $\{q_{0},q_{1},\dots,q_{L}\}$
$P$	Special modulus chain $\{q_{L+1},q_{L+2},\dots\}$
$L$	Length of the modulus chain (maximum multiplicative depth)
$l\in[0,L)$	Current remaining multiplicative level

Table 1. Notations and symbols used in this paper

Figure 1. Illustration of data layout and parallelization models in CKKS computation. (a) Packing that maps plaintext tensors into ciphertext slots. (b) Coefficient-parallel, distributing work across polynomial coefficients. (c) Limb-parallel, distributing work across RNS primes and exposing thread parallel on coefficient. Three-part illustration of CKKS data layout and parallelization. It shows plaintext values packed into ciphertext slots, coefficient-parallel execution across polynomial coefficients, and limb-parallel execution across RNS primes.

Figure 2. Encrypted Transformer inference pipeline with bootstrapping placement (NEXUS ( Zhang et al., 2025 ) ). Pipeline diagram of encrypted Transformer inference, showing linear layers, attention, nonlinear layers, and where bootstrapping is inserted to restore multiplicative depth.

Figure 3. Fundamental parallelisms in plaintext inference. (a) Data parallelism replicates model weights across devices. (b,c) Tensor parallelism partitions weights and aggregates partial results via collectives. Three-panel overview of plaintext inference parallelism. It contrasts data parallelism, which replicates weights, with tensor parallelism, which shards weights and then uses collectives to combine partial results.

Figure 4. Memory breakdown of encrypted BERT-Base inference across input sequence lengths. The stacked bars represent the major memory components. Common NVIDIA GPU memory capacities are marked for reference. Stacked bar chart of encrypted BERT-Base memory versus sequence length. Activation memory grows rapidly and becomes dominant at long sequences, exceeding common GPU memory capacities.

Collective	Plaintext cost	Plaintext example	Encrypted cost	Encrypted example
P2P send $(T$ tokens $)$	$T\cdot d_{\text{model}}\cdot\texttt{sizeof(bf16)}$	$1024\cdot 1024\cdot 2\approx 2$ MB	$K_{\text{send}}(T)\cdot(2\cdot l\cdot 2N\cdot 8)$	$32\times 40$ MB $\approx 1.3$ GB
AllReduce $(T$ tokens $)$	$(G-1)\,T\cdot d_{\text{model}}\cdot\texttt{sizeof(bf16)}$	$3\times 2$ MB $=6$ MB	$(G-1)\,K_{\text{send}}(T)\cdot(2\cdot l\cdot 2N\cdot 8)$	$3\times 1.3$ GB $\approx 3.8$ GB
AllGather $(T$ tokens $)$	$(G-1)\,\frac{T}{G}\cdot d_{\text{model}}\cdot\texttt{sizeof(bf16)}$	$3\times 0.5$ MB $=1.5$ MB	$(G-1)\,K_{\text{send}}\!\left(\frac{T}{G}\right)\cdot(2\cdot l\cdot 2N\cdot 8)$	$3\times(8\times 40$ MB $)\approx 0.96$ GB

Table 2. Per-device communication cost of common collectives in plaintext and encrypted inference. Example assumes

T{=}1024

d_{\text{model}}{=}1024

, bf16, full slot utilization,

S{=}32768

, and bootstrappable CKKS parameters (

\sim

40 MB ciphertexts at

l{=}35

). Encrypted communication is 2–3 orders of magnitude more expensive than plaintext.

Figure 5. (a) Collective in RNS-parallel BSGS matrix multiplication down to compute graph level. (b) Latency breakdown of ciphertext matrix multiplication on a 2-GPU with NVLink. Two-panel figure. Panel a shows collectives inserted repeatedly during RNS-parallel BSGS ciphertext matrix multiplication. Panel b shows that on two GPUs with NVLink, communication and key switching dominate ciphertext matrix multiplication latency.

Figure 6. (a) Encrypted diagonal matrix multiplication. (b) Directly applying plaintext tensor parallelism induces redundant ciphertext replication. Two-panel diagram. Panel a illustrates encrypted diagonal matrix multiplication. Panel b shows that directly reusing plaintext tensor parallelism replicates ciphertexts across devices and causes redundant encrypted work.

Figure 7. Output-slot coupling: (a) weight slices within the same output group generate results on the same slots; (b) splitting token output-coupled weight onto 2 devices leads to half slot waste in the next computation. Two-panel diagram of output-slot coupling in encrypted linear layers. Weight slices for the same output group map to the same slots, and splitting them across two devices leaves many output slots unused in the next step.

Figure 8. Parallel self-attention: (a) data partition and communication insertion; (b) latency hiding via operator reordering. Two-panel wide figure for encrypted self-attention. Panel a shows how data is partitioned and where collectives are inserted. Panel b shows operator reordering that overlaps communication with computation to hide latency.

Figure 9. The AEGIS compile pass. Block diagram of the \methodnamecompile pass, showing lowering from a high-level application representation to encrypted operators, then to polynomial-level instructions, followed by scheduling and multi-GPU execution planning.

Figure 10. Weak scaling: (a) upper bound on the number of devices as a function of input sequence length (assuming full slot utilization); (b) per-device memory consumption under the maximal feasible device count for each sequence length. Two-panel weak-scaling figure. Panel a plots the maximum usable device count against input sequence length under full slot utilization. Panel b plots per-device memory at that maximal device count.

Layer	Depth on Entry	Layer Latency Proportion (%)	Plaintext Tensor Parallel NVLink	HEBooster Limb Parallel NVLink	Cinnamon Limb Parallel NVLink	Hydra Tensor Parallel NVLink	AEGIS Parallel NVLink
$Q,K,V$ Linear Projection	0	28.78%	49.52%	31.71%	81.71%	89.53%	93.29%
$A\leftarrow Q\times K^{\intercal}/\sqrt{d}$	1	16.20%	50.00%	30.75%	30.75%	91.13%	92.37%
$A\leftarrow\texttt{Softmax}(A)$	2	3.48%	50.00%	30.43%	30.43%	40.52%	89.32%
$X\leftarrow A\times V$	18	3.29%	50.00%	30.75%	30.75%	91.13%	94.21%
$X\leftarrow X\times W_{o}$	19	1.20%	49.87%	34.63%	84.63%	88.32%	94.66%
Bootstrapping	20	1.14%	50.00%	31.50%	61.23%	48.63%	91.32%

Table 3. Scaling efficiency in terms of latency with different parallel strategies for a fixed 128 input tokens on 2 GPUs. Layer Latency Proportion is a reference value on single device execution.

Figure 11. Communication volume on two GPUs with 128 input tokens on BERT. Bar chart comparing communication volume in megabytes for different methods on two GPUs with 128-token BERT inference. The figure shows much higher communication for limb-parallel baselines than for \methodname.

Figure 12. End-to-end inference (a) per-device memory consumption; (b) full-model speed-up, where a baseline

1.0

corresponds to the latency of processing a full request on a single device. Hatched bar segments and hollow markers denote OOM on 48GB GPU. Two-panel figure for two-GPU end-to-end inference. Panel a shows per-device memory across input lengths for each method. Panel b shows full-model speed-up relative to single-GPU execution, with hatched or hollow marks indicating out-of-memory cases on a 48 GB GPU.

Figure 13. End-to-end inference on four GPUs: (a) per-device memory consumption; (b) full-model speed-up. Hatched entries denote OOM configurations with inferred values used for visualization. Two-panel figure for four-GPU end-to-end inference. Panel a shows per-device memory across input lengths for each method. Panel b shows full-model speed-up, with hatched entries marking out-of-memory configurations whose values are inferred for visualization.

Efficient Transformer Inference

Efficient Transformer Inference studies methods and systems that reduce the latency, memory, and energy needed to run Transformer models.

See all 121 papers

Cerium: A Multi-GPU Framework for Terabyte-Scale Encrypted Inference
Dec 12, 2025Siddharth Jayashankar, Joshua Kim, Michael B. Sullivan +2Homomorphic Encryption GPU Kernel Optimization
Carnegie Mellon University · UT Austin · NVIDIA
Layer-Parallel Inference Reduces Encrypted Nonlinear Depth in Transformers
Jul 6, 2026Ligong Han, Kai Xu, Hao Wang +3Efficient Transformer Inference Transformer Inference
MBZUAI IFM · Red Hat AI Innovation · MIT-IBM Watson AI Lab +3
ATLAS: Automated Approximation of Transformers for Efficient Homomorphic Inference in One Hour
Jul 26, 2026Jianhang Xie, Sicheng Tan, Vishnu Naresh Boddeti +1Efficient Transformer Inference Homomorphic Encryption
City University of Hong Kong · Shandong University · Michigan State University

Dec 12, 2025J/K move · Enter open · S save

Cerium: A Multi-GPU Framework for Terabyte-Scale Encrypted Inference

Siddharth Jayashankar, Joshua Kim, Michael B. Sullivan, Wenting Zheng+1

Carnegie Mellon University · UT Austin · NVIDIA

Homomorphic Encryption GPU Kernel Optimization GPU Acceleration Memory-Efficient Inference LLM Inference Acceleration

Open paper AskPDF arXiv

Encrypted AI using fully homomorphic encryption (FHE) enables inference directly over encrypted queries, providing strong privacy guarantees. However, its computational and memory overheads have limited practical deployment. Custom FHE accelerators improve performance, but rely on advanced manufacturing technologies that limit their accessibility. GPUs offer a more widely available alternative, yet achieving ASIC-class performance on GPUs is challenging. Large models such as LLMs compound these challenges by requiring optimized kernels, terabyte-scale memory management, and efficient execution across multiple devices. We present Cerium, a multi-GPU framework for large-scale FHE inference. Cerium combines a domain-specific language, an optimizing compiler, and a runtime system to automatically generate optimized GPU kernels, compress plaintext data, manage terabyte-scale memory footprints, and optimize multi-GPU communication. Implemented on NVIDIA GPUs, Cerium outperforms expert-written GPU libraries by up to 2.25x and achieves performance competitive with state-of-the-art FHE ASICs, matching the CraterLake ASIC. It is the first GPU system to execute CKKS bootstrapping in under 10 milliseconds, achieving 7.5 milliseconds, and is the first to demonstrate encrypted inference for BERT-Base and Llama3-8B prefill in 8 seconds and 43 seconds, respectively. The Cerium framework is open-source and available at https://github.com/CMU-CAOS/Cerium

Abstract

Efficient Transformer Inference

Explore similar work

Cerium: A Multi-GPU Framework for Terabyte-Scale Encrypted Inference

Layer-Parallel Inference Reduces Encrypted Nonlinear Depth in Transformers

ATLAS: Automated Approximation of Transformers for Efficient Homomorphic Inference in One Hour