PyramidFFT: Rearchitecting FFT with Matrix-Aligned Nested-Radix for Hierarchical Scratchpad Memory on AI Accelerators

Emerging AI accelerators equipped with large-capacity scratchpad memory (SPM) provide high computational throughput and large on-chip storage capacity. However, for memory-bound HPC algorithms such as Fast Fourier Transform (FFT), the hierarchical SPM organization cannot be fully exploited, leading to underutilized compute resources and exacerbated bandwidth bottlenecks. In this paper, we propose PyramidFFT, a memory-efficient FFT system co-designed for AI accelerators with large-capacity SPM. PyramidFFT leverages large-capacity SPM to address the memory-bound problem of FFT through three key techniques: Memory-Aware Nested Network Design tightly couples FFT computation with the memory hierarchy to reduce memory access frequency; Matrix-Aligned Butterfly Layout Optimization, which dynamically adjusts butterfly parallelism and transforms the FFT data layout to better utilize Tensor Core Units (TCUs); Symmetry-Driven Real FFT Compression, which removes conjugate symmetric components in real FFT to further ease memory pressure. The results demonstrate that PyramidFFT mitigates memory-bound limitations, achieving up to an 8.2× efficiency improvement over the baseline without memory-hierarchy mapping. By aligning with the structural characteristics of large-capacity SPM, PyramidFFT bridges the gap between FFT’s memory-bound behavior and the compute-dense design of modern AI accelerators, demonstrating the effectiveness of architecture-aware FFT co-design on SPM-based systems.