fhe-ntt

fhe-ntt is a batched, RNS, negacyclic Number-Theoretic Transform for homomorphic encryption and lattice cryptography, loadable through kernels. The reference baselines are schoolbook negacyclic convolution, compared exactly, and CRT reconstruction across the RNS base against a bignum convolution.

The NTT is the finite-field FFT: it turns polynomial multiplication in the RLWE ring Z_q[X]/(X^n + 1) into transform, pointwise product, inverse transform, and it is the dominant cost of every homomorphic multiply, rotation, and bootstrap in BGV, BFV, and CKKS. FHE runs in a residue number system, a large modulus split over word-size primes, so one operation is a batch of independent NTTs. This kernel runs that whole batch on the GPU, millions of transforms per second, with results that are exact integers rather than approximations.

The Cooley-Tukey stages of one lane, butterflies doubling in span each stage, over the eight-prime RNS stack; the roundtrip INTT(NTT(a)) == a returns bitwise, the negacyclic product equals schoolbook exactly, and the batch sustains 1.6 million transforms per second at n = 4096 over 8 primes.

Usage

import torch
from kernels import get_kernel

fhe = get_kernel("phanerozoic/fhe-ntt", version=1, trust_remote_code=True)

n = 4096
primes = [576460752308273153, 576460752312401921]   # 59-bit, each = 1 (mod 2n)
ctx = fhe.NTTContext(n, primes)                       # precompute twiddles once

# coefficients are [L, n] or [B, L, n] int64 in [0, q_l)
a = torch.stack([torch.randint(0, q, (n,), dtype=torch.int64, device="cuda") for q in primes])
b = torch.stack([torch.randint(0, q, (n,), dtype=torch.int64, device="cuda") for q in primes])

c = ctx.poly_mul(a, b)          # negacyclic product per prime, coefficient domain
A = ctx.ntt(a.clone())          # forward NTT in place
ctx.intt(A)                     # inverse, recovers a

version selects the release branch; trust_remote_code is required by kernels for publishers without the trusted-publisher mark. poly_mul(a, b, primes) is also exposed as a free function that caches a context for the shape.

API

Method

The forward transform is a Cooley-Tukey decimation-in-time NTT and the inverse a Gentleman-Sande decimation-in-frequency NTT, both with the negacyclic twist folded into the twiddles (a primitive 2n-th root of unity per prime), so no separate pre/post weighting is needed. Each butterfly is a Shoup modular multiply: with the twiddle w fixed, its reciprocal floor(w * 2^64 / q) is precomputed, and x * w mod q is one multiply-high and one conditional subtraction. Coefficients and twiddles are < q < 2^62 and carried in int64; the pointwise product, where both operands vary, uses a 128-bit multiply and one reduction. The whole [B, L, n] batch runs in log2(n) stage launches per transform.

Measured

Forward NTT, batch of 16 polynomials over 8 primes, transforms per second:

The transform is memory-bound integer work with no tensor-core dependency, so it runs on every architecture from sm_80 to sm_120.

Correctness

• Roundtrip identity: INTT(NTT(a)) == a bitwise for n from 256 to 8192 and primes from 14 to 59 bits.
• Polynomial multiply: poly_mul(a, b) equals the schoolbook negacyclic convolution in Z_q[X]/(X^n + 1) exactly, per prime, at 59-bit primes.
• RNS consistency: across an RNS base the per-prime results CRT-reconstruct to the exact integer negacyclic product (a ~2^175 modulus reproducing the full bignum convolution), so the independent single-prime transforms are jointly correct.
• Deterministic: integer arithmetic; repeated runs are bitwise identical.

Requirements and limits

• NVIDIA GPU with compute capability 8.0+.
• n a power of two; primes < 2^62, each = 1 (mod 2n) so the 2n-th root exists; coefficients int64 in [0, q).
• The transform is per-prime and per-polynomial independent; supply the RNS base as the L axis and batch polynomials on the B axis.

References

Longa and Naehrig, "Speeding up the Number Theoretic Transform for Faster Ideal Lattice-Based Cryptography" (2016); Harvey, "Faster arithmetic for number-theoretic transforms" (2014); the negacyclic RNS NTT of Microsoft SEAL and OpenFHE.

License

Apache-2.0.