Tutorial: Staged NTT for Large Transforms

This tutorial walks you through building a batched staged Number Theoretic Transform round-trip at N=16384 using cuPQC-NTT. You'll learn how SubSize<M> decomposes one large transform into two kernel passes, how to write the stage kernels for both directions, and how to size the shared memory and launch geometry each pass needs.

The NTT is the modular-arithmetic analogue of the FFT: it turns polynomial multiplication into cheap pointwise multiplication, which is why it sits on the critical path of lattice-based post-quantum schemes, fully homomorphic encryption, and zero-knowledge proof systems. A single-pass transform keeps its entire workspace in one block's shared memory, so it stops scaling once N grows past what a block can hold. The staged API solves exactly that problem: with SubSize<M> and K = N / M, each pass bounds its per-block working set to max(K, M) elements rather than the full N, so the transform no longer has to fit in a single block at all.

cuPQC-NTT requires cuPQC SDK 0.6.0 or newer, the release that introduced it. cuPQC 0.6 targets compute capabilities 8.0, 8.6, 8.7, 8.9, and 9.0. The library implements the cyclic transform, meaning arithmetic in the ring modulo x^N - 1.

Step 1: Project Setup

Clone the cuPQC repository:

git clone https://github.com/NVIDIA/cuPQC.git
cd cuPQC/examples/ntt

This will download all examples including the staged NTT example. The Makefile in this directory will compile all examples, including example_staged_ntt.

Step 2: Include Required Headers

Start by including the necessary cuPQC SDK headers:

#include <cstdlib>
#include <vector>
#include <iostream>

#include <ntt.hpp>

using namespace cupqc;

The ntt.hpp header provides the NTT descriptor operators, the staged transform entry points, the built-in prime fields, and the precomputed roots of unity.

Step 3: Understand the Staged Decomposition

Choose N and the sub-size M up front, since both go into the descriptor as compile-time constants:

constexpr uint32_t NTT_N = 16384;
constexpr uint32_t NTT_M = 256;
constexpr uint32_t NTT_K = NTT_N / NTT_M;  // 64

SubSize<M> splits an N-point NTT into two successive kernel passes with K = N / M:

Forward stage 1 runs M blocks, each performing a K-point transform over contiguous elements.
Forward stage 2 runs K blocks, each performing an M-point transform over strided elements.

The inverse reverses the stage ordering — its stage 1 is the M-point strided pass and its stage 2 is the K-point contiguous pass — and applies the 1/N normalization in stage 2. Neither pass ever needs more than max(K, M) elements of shared memory, which for N=16384 with M=256 means 256 elements instead of 16384.

Staged transforms cover 2^14 through 2^24 points, and the allowed SubSize<M> values narrow as N grows, so consult the feature documentation before committing to a pair. This example uses N = 16384 = 2¹⁴, split as M = 256, K = 64.

Step 4: Define the Staged NTT Types

Define one descriptor per direction. The staged behavior comes entirely from adding SubSize<M>() to an otherwise ordinary block-execution NTT descriptor:

using ForwardNTT = decltype(Algorithm<algorithm::NTT>()
                            + Direction<nttDirection::FORWARD>()
                            + Precision<uint32_t>()
                            + Size<NTT_N>()
                            + SubSize<NTT_M>()
                            + Block()
                            + BlockDim<128>());

using InverseNTT = decltype(Algorithm<algorithm::NTT>()
                            + Direction<nttDirection::INVERSE>()
                            + Precision<uint32_t>()
                            + Size<NTT_N>()
                            + SubSize<NTT_M>()
                            + Block()
                            + BlockDim<128>());

Algorithm<algorithm::NTT>() selects the transform.
Direction<nttDirection::FORWARD>() and Direction<nttDirection::INVERSE>() pick the direction; the two descriptors are otherwise identical.
Precision<uint32_t>() selects the coefficient and modulus width. uint16_t, uint32_t, and uint64_t are supported.
Size<NTT_N>() sets the transform length.
SubSize<NTT_M>() requests the staged decomposition with sub-size M, which is what makes the stage entry points available.
Block() with BlockDim<128>() requests block execution with 128 threads per block.

The descriptors expose the constants used later on: ForwardNTT::Size is N and ForwardNTT::BlockDim is the thread count, so the launch code never has to repeat them.

Both examples in this directory use KoalaBear, p = 2^31 - 2^24 + 1, which supports NTT sizes up to 2²⁴:

constexpr uint32_t p = cupqc::KoalaBear;                    // KoalaBear prime
constexpr uint32_t g = cupqc::KoalaBear_primitive_root_14;  // primitive 2^14-th root of unity mod p

const uint32_t g_inv = host_modpow(g, p - 2, p);  // g^{-1} mod p
const uint32_t N_inv = 2130576385;                // (2^14)^{-1} mod p

KoalaBear_primitive_root_14 is one of the precomputed *_primitive_root_S constants for transform size 2^S. The inverse root is derived on the host with a small modular-exponentiation helper (host_modpow) using Fermat's little theorem, and N_inv is the 1/N factor the inverse transform needs.

Step 5: Generate the Twiddle Tables

Twiddle generation is identical to the single-pass transform — the staged decomposition reuses the same tables across both of its passes:

template<class NTT>
__global__ void make_twiddles_kernel(uint32_t* twiddles, const uint32_t p, const uint32_t g) {
    NTT().make_twiddles(twiddles, p, g);
}

template<class NTT>
__global__ void transform_twiddles_to_mont_kernel(uint32_t* twiddles, const uint32_t p) {
    NTT().transform_twiddles_to_mont(twiddles, p);
}

Generate one table per direction and convert both into the Montgomery domain. This happens once and the tables are then reused for the entire batch:

const nttConst<uint32_t> scheme_const(p);

make_twiddles_kernel<ForwardNTT><<<1, 1>>>(d_twiddles,     p, g);
make_twiddles_kernel<InverseNTT><<<1, 1>>>(d_inv_twiddles, p, g_inv);
transform_twiddles_to_mont_kernel<ForwardNTT><<<1, ForwardNTT::BlockDim>>>(d_twiddles, p);
transform_twiddles_to_mont_kernel<InverseNTT><<<1, InverseNTT::BlockDim>>>(d_inv_twiddles, p);

The forward table is built from g, the inverse table from g^-1 mod p. nttConst<uint32_t> bundles the modulus and its derived Montgomery constants; it is constructed on the host from p and passed to the kernels by value.

Step 6: Write the Forward Stage Kernels

Each stage is its own kernel, and each follows the same load, execute, store pattern as the single-pass transform, with stage-qualified entry points. Both kernels take the polynomial base pointer and a batch_id, and use extern __shared__ so the workspace size can be chosen per stage at launch time:

__global__ void fwd_stage_1_kernel(uint32_t* data, const uint32_t* twiddles,
                                   const nttConst<uint32_t> scheme_const, const int batch_id) {
    extern __shared__ uint32_t sdata[];
    ForwardNTT().stage_1_load_to_mont(sdata, data + batch_id * ForwardNTT::Size, blockIdx.x, scheme_const);
    __syncthreads();
    ForwardNTT().stage_1_execute(sdata, twiddles, scheme_const.p);
    __syncthreads();
    ForwardNTT().stage_1_store(sdata, data + batch_id * ForwardNTT::Size, blockIdx.x);
}

__global__ void fwd_stage_2_kernel(uint32_t* data, const uint32_t* twiddles,
                                   const nttConst<uint32_t> scheme_const, const int batch_id) {
    extern __shared__ uint32_t sdata[];
    ForwardNTT().stage_2_load(sdata, data + batch_id * ForwardNTT::Size, blockIdx.x);
    __syncthreads();
    ForwardNTT().stage_2_execute(sdata, twiddles, scheme_const.p);
    __syncthreads();
    ForwardNTT().stage_2_store_from_mont(sdata, data + batch_id * ForwardNTT::Size, blockIdx.x, scheme_const);
}

Device transforms operate on Montgomery-domain data, so conversion happens at the memory boundary rather than per operation. Notice where that boundary falls across the two passes: stage 1 converts on the way in with stage_1_load_to_mont but stores with plain stage_1_store, and stage 2 reads with plain stage_2_load and only converts back out with stage_2_store_from_mont. The intermediate result written to global memory between the passes therefore stays in the Montgomery domain, and the round trip through global memory costs no conversions.

Every stage entry point takes blockIdx.x explicitly. That is how a block learns which sub-transform it owns, and it is what lets the same pointer serve M blocks in stage 1 and K blocks in stage 2. Both stages operate on the same polynomial pointer, and stage 2 reads the output that stage 1 wrote, so the kernel launch order is the dependency.

Step 7: Write the Inverse Stage Kernels

The inverse kernels mirror the forward ones, with one addition: stage_2_execute takes N_inv and applies the 1/N normalization as part of the final pass.

__global__ void inv_stage_1_kernel(uint32_t* data, const uint32_t* inv_twiddles,
                                   const nttConst<uint32_t> scheme_const, const int batch_id) {
    extern __shared__ uint32_t sdata[];
    InverseNTT().stage_1_load_to_mont(sdata, data + batch_id * InverseNTT::Size, blockIdx.x, scheme_const);
    __syncthreads();
    InverseNTT().stage_1_execute(sdata, inv_twiddles, scheme_const.p);
    __syncthreads();
    InverseNTT().stage_1_store(sdata, data + batch_id * InverseNTT::Size, blockIdx.x);
}

__global__ void inv_stage_2_kernel(uint32_t* data, const uint32_t* inv_twiddles,
                                   const nttConst<uint32_t> scheme_const, const uint32_t N_inv,
                                   const int batch_id) {
    extern __shared__ uint32_t sdata[];
    InverseNTT().stage_2_load(sdata, data + batch_id * InverseNTT::Size, blockIdx.x);
    __syncthreads();
    InverseNTT().stage_2_execute(sdata, inv_twiddles, scheme_const.p, N_inv);
    __syncthreads();
    InverseNTT().stage_2_store_from_mont(sdata, data + batch_id * InverseNTT::Size, blockIdx.x, scheme_const);
}

The stage numbering is per direction, not global. InverseNTT's stage 1 is the M-point strided pass and its stage 2 is the K-point contiguous pass, which is the reverse of the forward ordering. The launch geometry in the next step reflects that.

Step 8: Size the Shared Memory and Launch the Passes

Each stage has its own shared-memory requirement, and there is a separate constexpr helper for each of the four passes. All of them take N, M, and the precision type:

constexpr uint32_t N = ForwardNTT::Size;
constexpr uint32_t M = NTT_M;
constexpr uint32_t K = NTT_K;

constexpr size_t fwd_s1_smem = fwd_stage_1_ntt_shared_workspace_size<N, M, uint32_t>();
constexpr size_t fwd_s2_smem = fwd_stage_2_ntt_shared_workspace_size<N, M, uint32_t>();
constexpr size_t inv_s1_smem = inv_stage_1_ntt_shared_workspace_size<N, M, uint32_t>();
constexpr size_t inv_s2_smem = inv_stage_2_ntt_shared_workspace_size<N, M, uint32_t>();

Then launch four kernels per polynomial. The grid dimension alternates between M and K, and the pointer offset selects the polynomial:

for (unsigned int batch_id = 0; batch_id < batch; batch_id++) {
    fwd_stage_1_kernel<<<M, ForwardNTT::BlockDim, fwd_s1_smem>>>(d_data, d_twiddles, scheme_const, batch_id);
    fwd_stage_2_kernel<<<K, ForwardNTT::BlockDim, fwd_s2_smem>>>(d_data, d_twiddles, scheme_const, batch_id);

    inv_stage_1_kernel<<<K, InverseNTT::BlockDim, inv_s1_smem>>>(d_data, d_inv_twiddles, scheme_const, batch_id);
    inv_stage_2_kernel<<<M, InverseNTT::BlockDim, inv_s2_smem>>>(d_data, d_inv_twiddles, scheme_const, N_inv, batch_id);
}

Reading down the grid dimensions gives the whole decomposition at a glance: M, K going forward and K, M coming back. Block dimension stays at BlockDim<128> throughout, independent of the grid, so the 128 threads of a block cooperate on whichever sub-transform that block owns.

In a real application the two forward stages would bracket pointwise operations on the NTT-domain coefficients before the inverse stages run.

Step 9: Verify the Round Trip

The host code fills the batch with values 0, 1, 2, ... reduced mod p, keeps a copy as the reference, runs the round trip on the device, and then compares every coefficient of every polynomial:

std::vector<uint32_t> h_data(N * batch);
for (unsigned int b = 0; b < batch; b++) {
    for (uint32_t i = 0; i < N; i++) {
        h_data[b * N + i] = (b * N + i) % p;
    }
}
const std::vector<uint32_t> reference = h_data;

// ... cudaMalloc, cudaMemcpy to device, twiddle setup, the four-kernel loop, cudaMemcpy back, cudaFree ...

unsigned int errors = 0;
for (size_t i = 0; i < reference.size(); i++) {
    if (h_data[i] != reference[i]) errors++;
}
if (errors == 0) {
    std::cout << "Round-trip OK: " << batch << " polynomial(s) of degree " << N
              << " (M=" << M << " K=" << K << ") verified." << std::endl;
    return true;
}
std::cout << "Round-trip FAILED: " << errors << " coefficient(s) did not match." << std::endl;
return false;

Because the transform is in place, a successful round trip leaves d_data holding exactly the input again, so INTT(NTT(f)) == f must hold for every coefficient. The count of mismatches is reported rather than just a pass/fail flag, which distinguishes a single bad coefficient from a wholesale failure. main propagates the result so the program exits non-zero on any mismatch:

const unsigned int batch = 4;
const bool ok = staged_ntt_round_trip(batch, p, g, g_inv, N_inv);

std::cout << "\n" << (ok ? "Example completed successfully.\n"
                         : "Example failed.\n");
return ok ? EXIT_SUCCESS : EXIT_FAILURE;

Step 10: Build and Run

The Makefile will build all examples in the folder. Run the staged NTT example:

make
./example_staged_ntt

The Makefile expects the cuPQC SDK at /usr/local/cupqc-sdk, or at a user-specified path set through the CUPQC_SDK_DIR environment variable. cuPQC-NTT is a static library, and it is linked with link-time optimization through -dlto -lcupqc-ntt.

The shipped library is built for a specific set of transform sizes and sub-sizes. A configuration the type system accepts but the library does not provide compiles cleanly and then fails to link, so a link error is the signal to check your Size<N> and SubSize<M> pair against the supported list rather than to suspect your kernels.

Step 11: Understanding the Output

Expected output:

================================================================
Staged Number Theoretic Transform (NTT) Example
================================================================

This example demonstrates a batched staged NTT round-trip using
cuPQC SDK. SubSize<M> decomposes one large N-point transform into
two kernel passes, so each pass only needs max(M, N/M) elements of
shared memory instead of the full N. This is what makes transforms
larger than a single block's shared memory possible.

Configuration: N=16384, M=256, K=64, KoalaBear field, block execution

Round-trip OK: 4 polynomial(s) of degree 16384 (M=256 K=64) verified.

Example completed successfully.

If any coefficient does not survive the round trip, the Round-trip OK line is replaced by Round-trip FAILED: <count> coefficient(s) did not match., the closing line becomes Example failed., and the program exits non-zero.

Customization Tips

Choose `SubSize<M>` for Your N

SubSize<M> is the one parameter that has no single right answer. The pair (N, M) fixes K = N / M, and the per-block shared-memory requirement follows max(K, M), so the most balanced split — M close to the square root of N — minimizes the peak workspace. N=16384 with M=256 gives K=64 and a peak of 256 elements:

constexpr uint32_t NTT_N = 16384;
constexpr uint32_t NTT_M = 256;
constexpr uint32_t NTT_K = NTT_N / NTT_M;  // 64

Balance is not the only consideration, though. M and K also set the two grid dimensions, so a very lopsided split leaves one of the passes with too few blocks to fill the GPU. And the larger your N, the fewer valid SubSize<M> values remain from the ranges listed in Step 3, so check the supported pairs before tuning. Remember that an unsupported pair fails at link time, not compile time.

Scale the Batch Size

The batch size is a plain host variable, so raising it needs no descriptor change:

const unsigned int batch = 4;

Twiddle tables are generated once and shared across the whole batch, so the added cost is the allocation and copy of N * batch coefficients plus four more kernel launches per polynomial. The loop launches those four kernels per batch_id sequentially, which is a good place to look first if you need more throughput from a large batch.

Change the `Precision`

Precision works the same way here as in the single-pass case, described in the forward and inverse NTT tutorial, with one extra place to keep in step: the staged workspace helpers are parameterized by the element type too, so forward_stage1_ntt_shared_workspace_size<N, M, T>() and its three siblings change along with the descriptors, the nttConst<>, and the host buffers.

Prefer the Single-Pass Transform for Small N

The staged API exists for transforms whose working set exceeds one block's shared memory. Below that threshold it only adds cost: two kernel launches instead of one per direction, and a round trip through global memory between the passes. If your N fits in a single block, use the single-pass transform, whose load_to_mont, execute, and store_from_mont calls live in one kernel — see the forward and inverse NTT tutorial for that variant. That the staged range starts at 2^14 is a reasonable indication of where the trade-off begins to pay off.

Learn More

Example Source Code - Complete staged NTT example implementation
cuPQC-NTT User Guide - Usage guide with examples
cuPQC-NTT API Reference - Complete API documentation
cuPQC-NTT Features - Supported transform sizes and staged decomposition rules