Hierarchical Precision and Recursion for Accelerating Symmetric Linear Solves on MXUs

TL;DR

This work tackles the bottleneck of solving symmetric positive-definite systems by introducing a nested recursive Cholesky framework that recurses on $POTRF$, $TRSM$, and $SYRK$ to expose a rich hierarchy of $GEMM$-dominated subproblems on MXUs. It combines a tree-structured mixed-precision strategy, assigning $FP16$ to large off-diagonal blocks while preserving $FP32$/$FP64$ accuracy near the diagonal, with a lightweight per-block quantization scheme to maintain stability. Implemented in Julia, the solver achieves portable performance across NVIDIA $H200$ and AMD $MI300X$ backends, delivering up to $5\times$ speedups over full-precision baselines and significantly higher accuracy than pure $FP16$ in many configurations. The approach demonstrates strong cross-architecture applicability and lays the groundwork for hardware-aware, mixed-precision dense solvers at scale with potential extensions to block-sparse and multi-GPU settings for exascale workloads.

Abstract

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.

Figures (11)

• Figure 1: Nested Recursive Cholesky Decomposition Tree
• Figure 2: Recursive data structure enabling layered mixed precision. Off-diagonal blocks are computed in FP16 for performance, while diagonal regions use FP32 and/or FP64 to preserve accuracy.
• Figure 3: Blockwise quantization and dequantization in recursive mixed-precision algorithm.
• Figure 4: Speedup of Recursive SYRK on NVIDIA H200. Relative performance of the recursive Symmetric Rank-k update compared to cuBLAS FP64. The pure FP16 implementation achieves up to 149x speedup, while layered mixed-precision configurations deliver between 19x and 27x acceleration.
• Figure 5: Speedup of Recursive TRSM on NVIDIA H200. Performance of the recursive Triangular Solve normalized to cuBLAS FP64. The pure FP16 implementation achieves up to 6x speedup, while layered mixed-precision configurations deliver up to 5.3x acceleration.