💬 NLP

Enabling AI ASICs for Zero Knowledge Proof

This paper introduces MORPH, a framework that reformulates Zero-Knowledge Proof kernels using a hardware-aware complexity model to enable efficient execution on AI ASICs like TPUs, achieving up to 10x throughput improvements in NTT operations through MXU-centric arithmetic and optimized dataflow layouts.

Original authors: Jianming Tong, Jingtian Dang, Simon Langowski, Tianhao Huang, Asra Ali, Jeremy Kun, Jevin Jiang, Srinivas Devadas, Tushar Krishna

Published 2026-04-21

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Jianming Tong, Jingtian Dang, Simon Langowski, Tianhao Huang, Asra Ali, Jeremy Kun, Jevin Jiang, Srinivas Devadas, Tushar Krishna

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: The "Secret Math" Problem

Imagine you are running a massive bank that needs to prove to its customers that a transaction happened correctly without revealing the customer's name or the amount of money. This is called a Zero-Knowledge Proof (ZKP).

To do this, the bank's computer (the "Prover") has to solve incredibly difficult math puzzles. Two specific puzzles take up almost all the time:

MSM (Multi-Scalar Multiplication): Like trying to multiply a giant list of numbers by a giant list of points on a map, over and over again.
NTT (Number-Theoretic Transform): Like taking a long song and instantly breaking it down into its individual musical notes (frequencies) to check if the melody is correct.

The Problem: These puzzles are so heavy that they take hours to solve on standard computers. This makes ZKPs too slow and expensive for everyday use (like verifying a photo or a blockchain transaction).

The Solution: The researchers built MORPH, a new system that lets powerful AI chips (specifically Google's TPUs) solve these puzzles much faster.

The Analogy: The Factory vs. The Assembly Line

To understand why this is hard, imagine two different types of factories:

The General Factory (GPUs/CPUs): This factory has many skilled workers who can do almost anything. They are great at handling complex, messy tasks where every step is different. They are flexible but not super fast at doing the same thing millions of times.
The AI Factory (TPUs): This factory is built for mass production. It has a giant, super-fast conveyor belt (the MXU) designed to do one specific thing: multiply huge grids of numbers (matrices) all at once. It is incredibly fast and energy-efficient, but it is rigid. If you try to give it a messy, irregular task, it gets stuck because it can't handle the "mess."

The Conflict: ZKP math is messy. It involves huge numbers (hundreds of bits long) and requires constant "rearranging" of data. The AI Factory (TPU) is great at the math part but terrible at the rearranging part. If you just plug the old ZKP code into the TPU, it runs 30 times slower because the workers spend all their time shuffling papers around instead of doing the math.

How MORPH Fixes It: The Two-Step Magic

The MORPH team realized they couldn't just force the old code onto the new machine. They had to redesign the workflow to fit the AI Factory's strengths. They did this in two ways:

1. The Arithmetic Fix: "The Translation Layer"

The Problem: ZKP math uses huge numbers (like 377 bits long). The TPU's conveyor belt only understands small numbers (8 or 32 bits). The old way of handling this was like trying to carry a 500-pound boulder by breaking it into tiny pebbles, carrying them one by one, and then gluing them back together. It's slow and requires a lot of "glue" (carry chains).

The MORPH Solution: They invented a Translation Layer (RNS Lazy Reduction).

Analogy: Instead of carrying the boulder, they put it in a special, oversized shipping container that fits perfectly on the conveyor belt.
They convert the huge, messy math problem into a format that looks like a giant grid of small numbers.
Now, the TPU's super-fast conveyor belt can process the whole grid in one go, without stopping to glue anything together.
Result: This turns a slow, step-by-step process into a massive, parallel sprint. It's 90 times faster for the math part.

2. The Dataflow Fix: "The Stationary Warehouse"

The Problem: The old ZKP algorithms require the computer to constantly stop, shuffle data around, and reorganize it (like a librarian constantly moving books to different shelves just to find the next one). On a TPU, this "shuffling" is the biggest bottleneck. It's like having a Ferrari that spends 90% of its time stuck in traffic.

The MORPH Solution: They redesigned the workflow to be Layout-Stationary.

Analogy: Imagine a warehouse where the boxes are pre-sorted and placed exactly where the robot arm needs them. The robot never has to move the boxes around; it just grabs them and works.
They created a new method called LS-PPG (for the MSM puzzle) and a 5-step NTT (for the frequency puzzle).
These methods ensure that once the data is loaded, it stays in the perfect position for the TPU to chew through it. No more shuffling, no more traffic jams.
Result: The TPU spends almost 100% of its time doing math, not moving data.

The Results: A Speed Demon

By combining these two fixes, MORPH turned the TPU into a ZKP supercomputer.

Speed: For the "frequency" puzzle (NTT), the TPU is now 10 times faster than the best existing systems (like NVIDIA GPUs).
Efficiency: It uses much less energy to do the same work.
Scalability: As the math gets harder (larger numbers), the TPU doesn't slow down as much as other chips do.

The Takeaway

Think of MORPH as the ultimate adapter. It took a machine built for AI (Google's TPU) and rewired it to be the world's best machine for "Secret Math" (ZKPs).

Before: Trying to run a Formula 1 car on a dirt road (Slow, bumpy, inefficient).
After: MORPH paved a perfect highway specifically for that car. Now, it's the fastest vehicle on the track.

This breakthrough means that in the future, we could verify complex computations (like private blockchain transactions or secure AI processing) in seconds instead of hours, making privacy-preserving technology practical for everyone.

1. Problem Statement

Zero-Knowledge Proofs (ZKPs) are critical for scalable blockchains and verifiable computation, but their adoption is hindered by the prohibitive cost and latency of the prover. The prover's runtime is dominated by two kernels:

Multi-Scalar Multiplication (MSM): Accounts for ~70% of latency.
Number-Theoretic Transform (NTT): Accounts for ~20–30% of latency.

While AI ASICs like Google's TPUs offer massive matrix throughput and superior energy efficiency compared to GPUs, existing ZKP algorithms (optimized for GPUs/FPGAs) perform poorly on TPUs. The paper identifies two primary reasons for this inefficiency:

Arithmetic Mismatch: ZKPs require modular arithmetic over large prime fields (256–753 bits), but TPUs natively support low precision (8/16/32-bit). Traditional algorithms rely on sequential carry propagation chains (e.g., Radix Montgomery reduction) which are inefficient on TPUs because they require fine-grained data shuffling that the hardware cannot parallelize effectively.
Dataflow Mismatch: State-of-the-art (SotA) algorithms involve frequent layout transformations (shuffles, transposes) to align data with the TPU's coarse-grained Vector Registers (VRegs). These transformations incur massive latency that is invisible to traditional complexity models (Big-O) but dominate execution time on AI ASICs.

2. Methodology: The MORPH Framework

The authors propose MORPH, the first framework to systematically reformulate ZKP kernels for AI ASICs. It introduces a new complexity model and applies optimizations at both the arithmetic and dataflow levels.

A. New Complexity Model: Big- $T$

The paper introduces Big- $T$ complexity, a hardware-aware metric that extends Big-O.

Limitation of Big-O: It ignores heterogeneous bottlenecks and the cost of data layout transformations.
Big- $T$ Definition: It measures the execution time as the maximum of:
1. The work divided by the parallelism of each heterogeneous unit ( $W_k / P_k$ ).
2. The latency of layout transformations and memory access ($Mem$).
Insight: Big- $T$ reveals that on TPUs, the bottleneck is often not compute throughput but the sequential shuffle-bound carry propagation and layout reorganization.

B. Arithmetic Optimization: MXU-Centric RNS Lazy Reduction

To eliminate sequential carry chains, MORPH reformulates modular arithmetic:

Extended RNS Field: Instead of performing arithmetic directly in the prime field $\mathbb{F}_M$ , MORPH maps elements into a larger, non-prime field $\mathbb{F}_Q$ that supports a Residue Number System (RNS).
Lazy Reduction: Modular multiplication is decomposed into independent 32-bit vector multiplications (no carry propagation).
Matrix-Oriented Reduction: The subsequent modular reduction (RNS reconstruction $\to$ mod $M$ $\to$ RNS decomposition) is transformed into a low-precision dense matrix multiplication (GEMM) and vectorized operations.
Result: This shifts the bottleneck from sequential carry chains (XLU-bound) to parallel matrix multiplication (MXU-bound), achieving up to 90 $\times$ speedup over traditional Radix Montgomery reduction.

C. Dataflow Optimization: Layout-Stationary Algorithms

MORPH redesigns dataflows to minimize data movement and layout changes:

For MSM (LS-PPG): The authors propose Layout-Stationary Pippenger (LS-PPG).
- It enforces a unified sharding and layout strategy across the "Presorting" and "Bucket Accumulation" phases.
- This eliminates inter-TPU communication and intra-TPU layout reorganization (shuffles/transposes).
- It fuses presorting with bucket accumulation, allowing points to be consumed immediately after generation without writing back to off-chip HBM.
For NTT (5-Step NTT):
- While a 3-step NTT exists, it still incurs high memory span for large problem sizes.
- MORPH introduces a 5-step NTT, which recursively partitions the row-wise NTT of the 3-step version into a 3-step structure over smaller factors ( $R_1, R_2, C$ ).
- This reduces the effective MXU span and parameter storage requirements while maintaining high MXU utilization.

3. Key Contributions

Big- $T$ Complexity: A novel metric that quantifies the impact of heterogeneous compute spans and layout transformation costs, providing a principled way to analyze ZKP designs on AI ASICs.
MXU-Centric RNS Lazy Reduction: A technique that converts high-precision modular arithmetic into dense low-precision GEMMs, eliminating carry chains and achieving up to 90 $\times$ speedup over SotA radix decomposed Montgomery reduction.
LS-PPG and 5-Step NTT: Dataflow algorithms that avoid inter-device communication and minimize intra-device layout reorganization.
- LS-PPG: Delivers up to 5.8 $\times$ speedup over SotA MSM on TPUs.
- 5-Step NTT: Reduces parameter storage overflow and improves throughput for large degrees.
First TPU-based ZKP Accelerator: The first framework to enable TPUs as a viable, high-throughput platform for ZKP kernels.

4. Experimental Results

The framework was implemented in JAX and evaluated on Google TPUv6e (scaled to match the power consumption of an NVIDIA V100).

Throughput Gains:
- NTT: MORPH achieves up to 10 $\times$ higher throughput than GZKP (SotA GPU implementation) for 753-bit NTTs.
- MSM: MORPH achieves comparable throughput (up to 1.2 $\times$ ) to GZKP for 753-bit MSM, despite the architectural challenges.
Scalability:
- Precision Scaling: As modulus size increases from 256-bit to 753-bit, GPU latency increases by 6–7 $\times$ due to longer carry chains. In contrast, TPU latency only increases by 1.3–3 $\times$ because RNS arithmetic maps efficiently to the MXU.
- Batch Size: Performance improves significantly with larger batch sizes (up to 128), fully utilizing the TPU's VRegs and MXU.
Ablation Study:
- Arithmetic optimization alone provided 50–90 $\times$ latency reduction for MSM.
- Dataflow optimization provided an additional 3.1 $\times$ speedup for MSM.

5. Significance

This work fundamentally shifts the paradigm for ZKP acceleration:

AI ASICs as SotA Platforms: It demonstrates that AI ASICs (specifically TPUs), previously thought unsuitable for cryptographic workloads due to precision and dataflow constraints, can actually outperform GPUs for high-precision ZKP kernels when properly reformulated.
Energy Efficiency: By leveraging the massive parallelism and energy efficiency of TPUs, MORPH offers a path toward more sustainable and scalable verifiable computation.
Generalizable Framework: The Big- $T$ model and the strategy of "layout-stationary" execution provide a blueprint for porting other irregular, high-precision algorithms to modern heterogeneous AI hardware.

In conclusion, MORPH bridges the gap between the mathematical requirements of ZKPs and the architectural strengths of AI ASICs, making TPUs a leading platform for high-precision NTT and a competitive option for MSM.