XK.jl: Composable and Portable Multi-GPU BLAS in Julia

Abstract

This paper introduces XK.BLAS: the BLAS module of the Julia package XK.jl. This module provides BLAS APIs that are portable across all three major GPU vendors (AMD, Intel, NVIDIA) for multi-device architectures. XK.BLAS is built on top of the XKRT tasking runtime systems and the XKBlas library. An XKBlas program is a sequence of task-generating routines with no explicit handling of memory (such as allocation and data movement) or scheduling decisions (such as target device selection). Memory is instead lazily allocated, migrated, and evicted, while tasks are automatically distributed across available devices. Tasks are composable, so that fine-grained dependencies are inferred from the memory regions they access, allowing concurrent data movement and kernel execution across different routines when the dataflow permits. We evaluate primitives of XK.BLAS and its use as a multi-GPU backend for the package Krylov.jl. We show up to 2.8 × speedup when scaling from 1 to 4 H100 GPUs. More importantly, XK.BLAS enables solving problems that would not otherwise fit on a single GPU, with no code changes at all.