Symmetric linear solves are fundamental to a wide range of scientific and engineering applications, from climate modeling and structural analysis to machine learning and optimization. These workloads often rely on Cholesky (POTRF) decomposition and its supporting operations, triangular solves (TRSM) and symmetric rank-k updates (SYRK), which together form the computational core for solving symmetric positive-definite systems. To accelerate these kernels, we present a portable, mixed-precision solver designed for Matrix Processing Units (MXUs), including NVIDIA Tensor Cores (H200) and AMD Matrix Cores (MI300X). Our algorithm builds on a nested recursive formulation in which Cholesky exposes parallelism through recursive decomposition of its TRSM and SYRK sub-problems. This structure yields a hierarchical recursion that maximizes GEMM throughput while enabling fine-grained control over numerical precision. We introduce a custom recursive data structure that assigns low-precision FP16 arithmetic to large off-diagonal blocks, while preserving high precision on diagonal blocks to ensure numerical stability. The solver is implemented in Julia, leveraging array programming, multiple dispatch, and dynamic type inference to enable seamless expression of mixed-precision computation. This design provides a high-level, hardware-agnostic interface while efficiently interfacing with low-level vendor libraries for backend portability. On H200, our recursive FP64 SYRK achieves a 14x speedup over cuBLAS, while mixed-precision delivers up to 27x speedup in SYRK and 5x in TRSM over full-precision baselines. This results in a 5x overall speedup for Cholesky versus cuSOLVER FP64, with 100x better accuracy than pure FP16 while retaining 88% of its peak speedup. Comparable performance and accuracy trends are observed on MI300X, demonstrating broad applicability across GPUs.
翻译:对称线性求解是气候建模、结构分析、机器学习与优化等广泛科学与工程应用的基础。此类计算通常依赖Cholesky(POTRF)分解及其支撑操作——三角求解(TRSM)与对称秩-k更新(SYRK),三者共同构成求解对称正定系统的计算核心。为加速这些核心函数,我们提出一种面向矩阵处理单元(MXU,包括NVIDIA Tensor Core H200与AMD Matrix Core MI300X)的可移植混合精度求解器。算法基于嵌套递归框架,其中Cholesky通过递归分解其TRSM与SYRK子问题暴露并行性。该结构形成分层递归,在最大化GEMM吞吐量的同时实现数值精度的细粒度控制。我们引入自定义递归数据结构:对大型非对角块采用低精度FP16运算,而对角块保留高精度以确保数值稳定性。求解器以Julia实现,利用数组编程、多重分派与动态类型推断,无缝表达混合精度计算。该设计提供高层级、硬件无关的接口,同时通过底层供应商库实现高效后端可移植性。在H200上,递归FP64 SYRK相较cuBLAS实现14倍加速;混合精度在SYRK上达成27倍、TRSM上达成5倍的加速比(相较于全精度基线)。这使得Cholesky相较cuSOLVER FP64整体加速5倍,同时精度较纯FP16提升100倍,且保留其88%的峰值加速性能。在MI300X上观测到类似的性能与精度趋势,证明了该方法在GPU间的广泛适用性。