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间的广泛适用性。

0
下载
关闭预览

相关内容

重磅!几何深度学习 新书,160页pdf阐述
专知会员服务
265+阅读 · 2021年4月29日
一文读懂线性回归、岭回归和Lasso回归
CSDN
34+阅读 · 2019年10月13日
线性回归:简单线性回归详解
专知
12+阅读 · 2018年3月10日
从点到线:逻辑回归到条件随机场
夕小瑶的卖萌屋
15+阅读 · 2017年7月22日
各种相似性度量及Python实现
机器学习算法与Python学习
11+阅读 · 2017年7月6日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
6+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
VIP会员
最新内容
综述 | 3D场景图:开放挑战与未来方向
专知会员服务
4+阅读 · 6月22日
21世纪的无人机战争
专知会员服务
4+阅读 · 6月22日
《量子技术的军事任务技术适配与利用》
专知会员服务
4+阅读 · 6月22日
美国从乌克兰无人机战争中学习经验
专知会员服务
7+阅读 · 6月21日
ICML 2026 | 面向视觉语言模型的语义鲁棒性认证
专知会员服务
5+阅读 · 6月21日
相关基金
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
6+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
Top
微信扫码咨询专知VIP会员