Enabling AI ASICs for Zero Knowledge Proof
Jianming Tong, Jingtian Dang, Simon Langowski, Tianhao Huang, Asra Ali + 4 more
TLDR
MORPH is a framework that reformulates ZKP kernels for AI ASICs like TPUs, significantly boosting performance for MSM and NTT operations.
Key contributions
- MORPH framework reformulates ZKP kernels to efficiently run on AI ASICs like TPUs.
- Introduces Big-T complexity, a hardware-aware model for ZKP bottlenecks on ASICs.
- Develops MXU-centric extended-RNS lazy reduction for high-precision modular arithmetic.
- Achieves up to 10x higher throughput for NTT and comparable for MSM on TPUs.
Why it matters
Zero-knowledge proof (ZKP) provers are computationally expensive, especially for multi-scalar multiplication (MSM) and number-theoretic transforms (NTT). This paper leverages the efficiency of AI ASICs like TPUs by redesigning ZKP kernels. It offers significant performance gains, making ZKPs more practical and scalable.
Original 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.
📬 Weekly AI Paper Digest
Get the top 10 AI/ML arXiv papers from the week — summarized, scored, and delivered to your inbox every Monday.