Enabling AI ASICs for Zero Knowledge Proof

Abstract: Zero-knowledge proof (ZKP) provers remain costly because multi-scalar multiplication (MSM) and number-theoretic transforms (NTTs) dominate runtime as they need significant computation. AI ASICs such as TPUs provide massive matrix throughput and SotA energy efficiency. We present MORPH, the first framework that reformulates ZKP kernels to match AI-ASIC execution. We introduce Big-T complexity, a hardware-aware complexity model that exposes heterogeneous bottlenecks and layout-transformation costs ignored by Big-O. Guided by this analysis, (1) at arithmetic level, MORPH develops an MXU-centric extended-RNS lazy reduction that converts high-precision modular arithmetic into dense low-precision GEMMs, eliminating all carry chains, and (2) at dataflow level, MORPH constructs a unified-sharding layout-stationary TPU Pippenger MSM and optimized 3/5-step NTT that avoid on-TPU shuffles to minimize costly memory reorganization. Implemented in JAX, MORPH enables TPUv6e8 to achieve up-to 10x higher throughput on NTT and comparable throughput on MSM than GZKP. Our code: https://github.com/EfficientPPML/MORPH.

• The paper introduces MORPH, a hardware-aware framework that repurposes AI ASICs to achieve up to 90× speedup in modular arithmetic for zero-knowledge proofs.
• It presents Big-T, a new complexity model that emphasizes memory bandwidth and dataflow over arithmetic operations in ZKP workloads.
• By co-optimizing arithmetic (MXU-centric RNS Lazy Reduction) and dataflow (layout-stationary kernels), the approach significantly boosts throughput and reduces latency.

Enabling AI ASICs for Zero Knowledge Proof: A Technical Essay

Introduction

Zero-knowledge proofs (ZKPs) are increasingly fundamental for privacy-preserving verification in distributed and blockchain systems, but their computational demands—especially those of the prover—remain prohibitive due to two kernel bottlenecks: multi-scalar multiplication (MSM) and number-theoretic transform (NTT). The paper "Enabling AI ASICs for Zero Knowledge Proof" (2604.17808) addresses these challenges by introducing MORPH, a framework that systematically repurposes AI application-specific integrated circuits (ASICs), particularly Google TPUs, for ZKP provers. The technical contributions involve hardware-guided algorithmic reformulation, a new complexity model (Big-$T$), and both arithmetic and dataflow-level innovations that co-optimize for AI ASIC architectures.

Hardware Bottlenecks in ZKP Kernels

Classical MSM and NTT algorithms, optimized for general-purpose hardware, are fundamentally mismatched with AI ASICs such as TPUs. These devices exhibit distinctive heterogeneous, deeply pipelined, and highly parallel compute resources, organized around coarse granularity vector registers (VRegs). State-of-the-art ZKP algorithms, when directly ported, are hamstrung by three bottlenecks:

1. Precision Mismatch: ZKPs require high-precision modular arithmetic over large primes (256–753 bits), but AI ASICs natively support <32-bit computations, leading to low resource utilization.
2. Layout-Induced Stalls: Algorithmic dataflows generate fine-grained memory accesses, causing excessive and costly data shuffles and transposes on architectures optimized for contiguous VReg access.
3. Communication Overheads: MSM protocols (e.g., Pippenger’s) and multi-step NTTs impose significant data layout transformations and inter-core communication, further reducing throughput.

The following schematic illustrates the critical linkage between ZKP prover kernels and underlying AI ASIC operations:

(Figure 1)

Figure 1: MORPH's deployment flow with dataflow and arithmetic optimizations to accelerate ZKP on TPU.

Big-$T$ Complexity: A Hardware-Aware Model

Traditional algorithmic complexity models (such as Big-$O$) fail to reflect the true performance bottlenecks on heterogeneous parallel systems, as they ignore platform-specific compute and memory access patterns. The authors introduce Big-$T$ complexity, modeling runtime as the maximum of pipeline spans across heterogeneous compute units and memory subsystem latency:

$$T(N) = O\left(\max \left\{ \max_{k} \frac{W_k}{P_k}, \;\text{Mem} \right\}\right)$$

Where $W_k$ and $P_k$ denote workload and parallelism for unit $k$, respectively. Big-$T$ reveals that layout transformations and memory bandwidth, rather than arithmetic ops alone, dominate the critical path in ZKP workloads on AI ASICs.

This analysis motivates targeting not only computational patterns but also data movement and storage layout in algorithmic redesign.

Arithmetic-Level Innovations: MXU-Centric RNS Lazy Reduction

The dominant bottleneck in classical MSM and NTT arithmetic is high-precision modular reduction, which on non-native platforms requires digit decomposition and extensive carry propagation. This process is both memory- and shuffle-bound.

MORPH's solution: reformulate high-precision prime field operations as 32-bit vector multiplications in a larger extended RNS field, encapsulating prime modular multiplication in an efficient carry-free form. Key aspects:

• Extended Residue Number System (RNS): Large prime field elements are represented in a larger, factorable field, permitting independent low-precision multiplications.
• MXU-centric Lazy Reduction: Modular reduction is structured as a series of batched low-precision matrix multiplies, eliminating carry propagation and mapping perfectly onto the matrix multiplication unit (MXU) in TPUs.
• Outcome: All modular reduction stages are efficiently parallelized, shifting the bottleneck from shuffles to throughput-bound vector and matrix operations.

The flow is illustrated as follows:

(Figure 2)

Figure 2: Illustration of actual value change and computational flow of MXU-centric RNS Lazy Modular Multiplication.

This results in up to 90$\times$ speedup over state-of-the-art radix Montgomery reduction schemes for ZKP arithmetic kernels.

Dataflow-Level Optimizations: Unified Layout-Stationary Kernels

For both MSM and NTT, the challenge is to restructure algorithms so that their most communication- and layout-intensive stages (sorting, bucketing, twiddle factor permutations) are minimized or eliminated.

MSM: Layout-Stationary Pippenger (LS-PPG)

• Sharding without Communication: By aligning sharding and processing phases to the native TPU data layout, MORPH eliminates inter-TPU and intra-TPU communication during the dominant bucket accumulation (BA) phase.
• On-Chip Data Reuse: Presorting is fused with BA such that once points are read from memory, all ensuing computation occurs on-chip and is layout-invariant.
• Outcome: Memory passes are minimized and MSM throughput is systematically improved.

The following figure clarifies the communication and layout-reducing dataflow:

(Figure 3)

Figure 3: Illustration of LS-PPG's computational phases, focusing on fusion of data sharding and accumulation to minimize memory moves and layout transformations.

NTT: Recursive 5-Step Layout-Invariant Decomposition

• From 3-Step to 5-Step NTT: The authors generalize the 3-step layout-invariant NTT method to a 5-step recursive formulation, partitioning the transformation into smaller-gemms that match the MXU’s native tile sizes and further reduce parameter storage demand.
• Zero Layout Reorganization: All computation is performed on tiles that are natively aligned with the hardware layout, and explicit shuffling is avoided.
• Outcome: Significantly improved scalability and throughput at high precision and degree.

This recursive layout-optimized NTT computation is visualized as:

(Figure 4)

Figure 4: Performance analysis and computational steps in different NTT approaches. (c) shows MORPH's proposed 5-step NTT recursively reducing the row-wise NTT to minimize overall computation and layout costs.

Quantitative Results

Ablation and Comparative Performance

Extensive hardware experiments on Google TPUv6e8 (matched to NVIDIA V100 in power) demonstrate the impact of these optimizations:

• Arithmetic Optimization: Yields 50–90$T$0 latency reduction for MSM, mainly via the carry-free RNS reduction.
• Dataflow Optimization: Delivers up to 3.1$T$1 speedup for MSM, with the largest improvement in bucket reduction due to increased parallelism.
• NTT Throughput: MORPH achieves up to 10$T$2 throughput over GZKP (GPU baseline) for 753-bit high-precision NTT, establishing TPUs as the new state-of-the-art for large field NTTs.

Performance scaling with increasing workload degree, as well as batch size, is shown in the following ablations:

Figure 5: MORPH ablation study under different polynomial degrees, highlighting progressive contributions of arithmetic and dataflow-level optimizations.

Figure 6: Batch size scaling for modular multiplication and NTT, demonstrating improved latency with increasing batch size and saturation at high batch counts.

Hardware Utilization

A deeper perspective on the TPU architecture clarifies the source of these gains:

Figure 7: TPU programming model. Compute units operate on VRegs of shape $T$3 32-bit registers. Peak compute efficiency requires layouts enabling direct tile loading from VRegs; otherwise, layout transformations create performance bottlenecks.

Theoretical and Practical Implications

Practical implications include the demonstration that general-purpose AI ASICs—such as those deployed for deep learning at scale—are viable and competitive platforms for privacy-preserving cryptography tasks when ZKP kernels are appropriately refactored. This reduces the need for purpose-built ASICs and accelerates the adoption of ZKPs in production systems, including blockchains and confidential cloud services.

Theoretical implications emerge from the introduction and validation of Big-$T$4 complexity, which stands to influence future algorithmic analysis and hardware/software co-design methodologies for a variety of large-scale, bandwidth-sensitive parallel workloads.

Future Directions

Prospective developments include:

• Automated Code-Generation: Leveraging hardware-aware compilers that can automatically refactor cryptographic kernels into Big-$T$5 optimal forms for new AI ASIC architectures.
• Dedicated Reordering Units: Designing new on-chip hardware to further reduce or hide latency of fine-grained scatter operations in MSM dataflows.
• Extension to Other Protocols: Applying the MORPH design methodology to protocols beyond SNARKs, such as STARKs or succinct non-interactive arguments for machine learning.

Conclusion

This work demonstrates that by blending hardware-aware algorithm reformulation (arithmetic and dataflow) with a new asymptotic complexity model, ZKP prover kernels can be efficiently mapped onto AI ASICs, removing legacy performance bottlenecks and establishing TPUs as a competitive platform for scalable, high-throughput, privacy-preserving computation (2604.17808). The MORPH framework’s results point to a broader trend: the emergence of hardware/software co-design as essential for next-generation cryptographic and AI workloads.