GPU Libraries
The layer between your model code and the hardware — BLAS, DL primitives, and kernel-generation frameworks.
TopNotchNote
Notes that matter
why frameworks lean on BLAS libraries
PyTorch/TensorFlow/JAX do not hand-write matrix multiply kernels — almost every dense linear-algebra op eventually calls into a vendor BLAS library. These libraries exist because GEMM performance is extremely sensitive to tiling, register blocking, and memory-access patterns that must be re-tuned per GPU generation; re-deriving that tuning in every framework would be enormous duplicated effort.
BLAS libraries
| Purpose | Nvidia | AMD |
|---|---|---|
| Dense linear algebra (GEMM etc.) | cuBLAS | rocBLAS |
| Lightweight/AI-focused GEMM, low-precision + epilogue fusion | cuBLASLt | hipBLASLt |
| Vendor-agnostic wrapper | hipBLAS (dispatches to cuBLAS under the hood) | hipBLAS (dispatches to rocBLAS under the hood) |
hipBLASLt (and cuBLASLt) exist specifically for AI workloads: they expose GEMM epilogues (bias-add, activation, gradient scaling fused into the same kernel) and heuristic/algorithm-search autotuning that a classic BLAS API doesn't — this is the layer that actually gets tuned per-shape for transformer workloads.
deep learning primitive libraries
| Purpose | Nvidia | AMD |
|---|---|---|
| Conv/attention/normalization primitives | cuDNN | MIOpen |
| Collective communication (multi-GPU) | NCCL | RCCL |
| FFT | cuFFT | rocFFT |
| Sparse linear algebra | cuSPARSE | hipSPARSE |
| Primitive/parallel algorithms (sort, scan, reduce) | Thrust / CUB | rocThrust / rocPRIM |
kernel-generation frameworks
Above the fixed-function BLAS layer sits a set of tools for generating custom fused kernels rather than calling a pre-built op:
| Framework | Vendor | What it is |
|---|---|---|
| CUTLASS | Nvidia | C++ template library for building custom GEMM/conv kernels at near-cuBLAS performance |
| Composable Kernel (CK) | AMD | AMD's analogue of CUTLASS — templated building blocks for custom fused kernels |
| Triton | vendor-agnostic | Python-embedded DSL that lowers through LLVM to both Nvidia (PTX, via LLVM's NVPTX backend) and AMD (GCN/RDNA ISA, via LLVM's AMDGPU backend) — not related to Composable Kernel, which is a separate templated kernel library — what torch.compile generates for fused kernels |
autotuning in GPU libraries
GEMM performance depends on tile size, split-K, and pipelining choices that vary by shape, dtype, and GPU model — no single kernel is fastest everywhere. hipBLASLt/cuBLASLt and Triton-based kernels ship many candidate kernel variants and pick one per shape via either an offline tuned lookup table or an online heuristic/brute-force search the first time a new shape is seen. This is why the first call with a new shape is sometimes slower than subsequent calls — it may be triggering a search.
related topics
CUDA & HIP — the CUDA/HIP layer these libraries are built on.
PyTorch on GPU — how PyTorch dispatches into libraries like this.
Machine Learning Notes — the ML workloads that depend on these libraries.
PyTorch on GPU — how PyTorch dispatches into libraries like this.
Machine Learning Notes — the ML workloads that depend on these libraries.