papers/zoo-key-management/sections/07-gpu-ntt-acceleration.tex

\section{GPU NTT Acceleration}


\subsection{Number Theoretic Transform}

Both ML-KEM and ML-DSA rely heavily on polynomial multiplication in $R_q = \mathbb{Z}_q[X]/(X^{256} + 1)$. The Number Theoretic Transform (NTT) converts polynomial multiplication from $O(n^2)$ to $O(n \log n)$:

\begin{equation}
\hat{f}[k] = \sum_{j=0}^{n-1} f[j] \cdot \psi^{(2 \cdot \text{brv}(k) + 1) \cdot j} \mod q
\end{equation}

where $\psi$ is a primitive $2n$-th root of unity modulo $q$ and $\text{brv}$ is the bit-reversal permutation.

\subsection{GPU Kernel Design}

\begin{proposition}[GPU NTT Speedup]
For batch sizes $B \geq 1024$ polynomial multiplications with $n = 256$ and $q = 8380417$ (ML-DSA modulus), GPU NTT achieves $14\times$ throughput improvement over optimized AVX2 CPU implementations.
\end{proposition}

The kernel uses coalesced memory loads with bit-reversal permutation, butterfly operations in shared memory ($\log_2 n = 8$ stages), branchless Barrett reduction, and coalesced output writes.

\subsection{Throughput Analysis}

\begin{table}[h]
\centering
\begin{tabular}{lcc}
\toprule
\textbf{Implementation} & \textbf{NTT/sec (single)} & \textbf{NTT/sec (batch 4096)} \\
\midrule
Reference C & 185,000 & 185,000 \\
AVX2 (x86\_64) & 1,200,000 & 1,200,000 \\
GPU (RTX 4090) & 890,000 & 16,800,000 \\
\bottomrule
\end{tabular}
\caption{NTT throughput by implementation}
\label{tab:ntt-throughput}
\end{table}

GPU acceleration is most effective for batch operations during epoch transitions (bulk key generation) and high-throughput signing periods.

\subsection{Formal GPU Correctness}

The GPU NTT kernel is verified in \texttt{GPU/ConsensusScaling.lean}, proving output equivalence with the reference implementation, constant-time execution (no input-dependent branching), and memory bounds safety for $n \leq 1024$.