Toward Capturing Genetic Epistasis From Multivariate Genome-Wide Association Studies Using Mixed-Precision Kernel Ridge Regression

TL;DR

This paper introduces a GPU-accelerated, mixed-precision kernel ridge regression approach for large-scale genome-wide association studies, significantly improving performance and enabling analysis of extensive genetic data.

Contribution

It presents a novel mixed-precision GPU implementation of KRR for GWAS, leveraging tensor cores and symmetry to accelerate computations and scale to larger datasets.

Findings

Achieved 1.805 ExaOp/s performance on GPU.

Outperformed CPU-based GWAS software by five orders of magnitude.

Enabled large-scale multivariate genetic epistasis analysis.

Abstract

We exploit the widening margin in tensor-core performance between [FP64/FP32/FP16/INT8,FP64/FP32/FP16/FP8/INT8] on NVIDIA [Ampere,Hopper] GPUs to boost the performance of output accuracy-preserving mixed-precision computation of Genome-Wide Association Studies (GWAS) of 305K patients from the UK BioBank, the largest-ever GWAS cohort studied for genetic epistasis using a multivariate approach. Tile-centric adaptive-precision linear algebraic techniques motivated by reducing data motion gain enhanced significance with low-precision GPU arithmetic. At the core of Kernel Ridge Regression (KRR) techniques for GWAS lie compute-bound cubic-complexity matrix operations that inhibit scaling to aspirational dimensions of the population, genotypes, and phenotypes. We accelerate KRR matrix generation by redesigning the computation for Euclidean distances to engage INT8 tensor cores while exploiting symmetry.We accelerate solution of the regularized KRR systems by deploying a new four-precision Cholesky-based solver, which, at 1.805 mixed-precision ExaOp/s on a nearly full Alps system, outperforms the state-of-the-art CPU-only REGENIE GWAS software by five orders of magnitude.