English

PolyKAN: Efficient Fused GPU Operators for Polynomial Kolmogorov-Arnold Network Variants

Distributed, Parallel, and Cluster Computing 2025-11-20 v1 Artificial Intelligence

Abstract

Kolmogorov-Arnold Networks (KANs) promise higher expressive capability and stronger interpretability than Multi-Layer Perceptron, particularly in the domain of AI for Science. However, practical adoption has been hindered by low GPU utilization of existing parallel implementations. To address this challenge, we present a GPU-accelerated operator library, named PolyKAN which is the first general open-source implementation of KAN and its variants. PolyKAN fuses the forward and backward passes of polynomial KAN layers into a concise set of optimized CUDA kernels. Four orthogonal techniques underpin the design: (i) \emph{lookup-table} with linear interpolation that replaces runtime expensive math-library functions; (ii) \emph{2D tiling} to expose thread-level parallelism with preserving memory locality; (iii) a \emph{two-stage reduction} scheme converting scattered atomic updates into a single controllable merge step; and (iv) \emph{coefficient-layout reordering} yielding unit-stride reads under the tiled schedule. Using a KAN variant, Chebyshev KAN, as a case-study, PolyKAN delivers 1.21.2--10×10\times faster inference and 1.41.4--12×12\times faster training than a Triton + cuBLAS baseline, with identical accuracy on speech, audio-enhancement, and tabular-regression workloads on both highend GPU and consumer-grade GPU.

Keywords

Cite

@article{arxiv.2511.14852,
  title  = {PolyKAN: Efficient Fused GPU Operators for Polynomial Kolmogorov-Arnold Network Variants},
  author = {Mingkun Yu and Heming Zhong and Dan Huang and Yutong Lu and Jiazhi Jiang},
  journal= {arXiv preprint arXiv:2511.14852},
  year   = {2025}
}
R2 v1 2026-07-01T07:44:07.806Z