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

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.

• The paper introduces AEGIS, a framework for scalable long-sequence encrypted transformer inference using hybrid parallelism on multi-GPU systems.
• It employs dependency-aware partitioning and operator reordering to align ciphertext operations, substantially reducing communication overhead and memory usage.
• Empirical results demonstrate up to 96.62% scaling efficiency, a 3.86× speedup, and 69.1% memory reduction compared to prior state-of-the-art strategies.

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

---

Introduction and Motivation

The proliferation of MLaaS for transformer-based LLMs has accentuated privacy risks inherent in client-cloud inference settings. Fully Homomorphic Encryption (FHE) schemes, in particular RLWE-based CKKS, provide robust cryptographic guarantees by enabling arithmetic directly on encrypted data. However, naive encrypted transformer inference rapidly exceeds single-GPU memory limits due to sizeable encoded weights and activations—especially as input sequence length grows. For example, encrypted BERT-Base inference with 2048 tokens demands over 700 GB for weights and 110 GB for activations, surpassing any commodity GPU's capacity and mandating distributed execution.

Existing parallelization approaches from plaintext LLMs—data parallelism, tensor parallelism, and their hybrids—fail to scale for encrypted workloads. This is primarily due to rigid constraints induced by ciphertext packing and RNS modulus decomposition, leading to intractable communication and slot redundancy. Parallelism in CKKS, such as limb- or coefficient-parallelism, is application-agnostic and often communication-bound, as demonstrated by frequent key switching and rotation chains. Even advanced multi-GPU HE accelerators exhibit superlinear slowdowns, e.g., 16.7s for matrix multiplication on one GPU vs. 62.3s on four GPUs, with communication dominating latency.

Figure 1: Illustration of CKKS computation layouts, packing plaintext tensors, coefficient-parallel, and limb-parallel execution across RNS primes.

The mismatch is further intensified by the extreme communication amplification: transmitting 1024 tokens requires only $\sim$2 MB in plaintext but expands to $\sim$1.3 GB under encryption. Thus, effective scaling is fundamentally constrained by encryption-level structure and application-level coupling.

---

Overview of the AEGIS Framework

AEGIS introduces a compiler-runtime co-design for multi-GPU encrypted transformer inference by orchestrating device placement based on ciphertext dependencies jointly induced by transformer dataflow semantics and CKKS polynomial coupling. The framework prioritizes:

• Modulus-coherent placement: Groups RNS limbs mandated for joint key switching or rotation on the same device, eliminating cross-device synchronization.
• Token-coherent placement: Aligns slices of the same token within devices, minimizing collectives necessitated by embedding-wise aggregation.
• Operator reordering: Restructures polynomial instruction schedules to overlap collective communication with computation, hiding latency deterministically.

AEGIS's compile pass performs high-level dependency analysis, inserts encrypted collectives at the polynomial level, and yields a deterministic execution plan for distributed encrypted inference.

Figure 2: Encrypted transformer inference pipeline, illustrating the placement of bootstrapping operations.

Figure 3: Parallelisms in plaintext inference, showing data and tensor parallelism, highlighting their limitations in encrypted domains.

---

Technical Contributions

Workload Partitioning

Dependency-aware partitioning in AEGIS ensures that polynomials necessary for joint reducibility—either due to modulus or token coupling—are co-located, avoiding slot underutilization and waste. In encrypted FFNs, AEGIS performs slot-wise accumulation, where output group slices are aggregated within the same device; splitting them leads to partial slot utilization and unnecessary communication.

Figure 4: Memory breakdown of encrypted BERT-Base inference across input sequence lengths; activations dominate at long sequences.

Figure 5: Applying plaintext tensor parallelism directly to homomorphic inference results in redundant ciphertext replication.

Self-Attention Partitioning

AEGIS distributes QKV slices and attention computation, applies ReduceScatter for distributed softmax, and leverages AllGather for result aggregation—fully utilizing token and modulus coherence while minimizing frequency and volume of collectives.

Figure 6: Output-slot coupling demonstrates the inefficiency of splitting token output-coupled weights across devices.

Figure 7: Parallel encrypted self-attention, with operator reordering for communication latency hiding.

Communication Operator Scheduling

AEGIS advances communication placement with three rules:

• Send before bootstrapping to minimize transfer size.
• Rescale prior to send reduces modulus depth and transfer volume.
• Local reduction before send optimizes data movement.

Operator reordering exposes communication opportunities earlier and overlaps collectives with computation, eliminating idle GPU streams and mitigating communication-induced stalls.

Figure 8: The AEGIS compile pass, illustrating dependency analysis, operator lowering, and execution planning.

---

Empirical Results

Weak Scaling and Memory Efficiency

AEGIS demonstrates scalability unconstrained by limb-count limitations or replication of ciphertexts. Per-device memory remains near optimal and avoids OOM across longer sequences, unlike tensor-parallel and hoisting-heavy baselines. Memory footprint reductions are significant: up to 69.1% lower per device compared to baselines.

Figure 9: Weak scaling, showing the maximum feasible device count and per-device memory for various sequence lengths.

Communication Overhead and Scaling Efficiency

AEGIS achieves substantial reductions in inter-GPU communication—up to 81.3% in self-attention and 57.9% in FFNs over prior SOTA. Operator reordering and dependency-aware placement together contribute up to 88.26% of speedup for long sequences. Layer-wise scaling efficiency consistently exceeds 90%, and end-to-end efficiency reaches 92.98% (two GPUs) and 96.62% (four GPUs).

Figure 10: Communication volume comparison for various strategies on encrypted BERT inference.

AEGIS consistently outperforms HEBooster, Cinnamon, Hydra, and Megatron-LM style parallelism in both memory and speedup metrics. On four GPUs, it achieves a 3.86$\times$ speedup, with latency scaling factors up to 0.95 even at 2048 tokens.

Figure 11: End-to-end inference scaling results; per-device memory and speed-up on two-GPU system.

Figure 12: End-to-end inference scaling on four GPUs; memory and speed-up, with OOM configurations inferred for completeness.

---

Implications and Future Directions

The results establish that scalable encrypted transformer inference demands coordinated application-encryption orchestration. Packing, placement, and communication scheduling must be solved as a cross-layer optimization problem; generic adaptation of plaintext parallelism is fundamentally inadequate due to rigid ciphertext packing and RNS coupling.

AEGIS's compiler-runtime approach is practically relevant for future privacy-preserving LLMs, especially in regulatory-bound settings. The decoupling of application and encryption dependencies enables scaling to ever-longer contexts, and provides a template for encrypted model serving, as well as system-level co-designs for future hardware accelerators.

There is scope for extending AEGIS to hybrid models and integrated distributed training, as well as to other RLWE-based schemes and logic/arithmetic mixed workloads. The deterministic communication hiding and dependency-aware placement may inform OS/runtime scheduling for encrypted AI at scale. As encrypted inference becomes more mainstream in MLaaS, coordinated orchestration will be pivotal for feasibility and performance.

---

Conclusion

AEGIS defines a new paradigm for scalable multi-GPU encrypted transformer inference, jointly orchestrating application-level data flow and encryption-level polynomial execution. It demonstrates superior scaling efficiency and memory reduction, with strong numerical results: up to 96.62% scaling efficiency, 3.86$\times$ speedup, and 69.1% memory improvement. The findings underscore the necessity of cross-layer optimization in encrypted AI systems, moving beyond direct adaptation of plaintext parallelism, and highlight directions for future development in encrypted model serving and system-level accelerator co-design.

---