Accelerating Frameworks
Getting the best performance out of an accelerator - whether one high-end GPU or a mixed fleet of GPUs, TPUs, and CPUs - is what ML frameworks are built for. Frameworks like PyTorch and JAX abstract away the hardware so you describe what to compute, not how to run it on a given chip. Underneath, technologies like CUDA kernels, the XLA compiler, and the PyTorch/XLA and JAX/XLA integrations turn that high-level code into fast, device-specific machine code.
Maximizing performance on specific hardware
On a single high-performance device like a GPU, the goal is to fully exploit its architecture: minimize idle cycles, maximize computational throughput, and optimize data movement. Frameworks such as TensorFlow, PyTorch, and CUDA hide the low-level complexity so researchers build and train models instead of managing hardware intricacies.
The three technologies that unlock this, from lowest to highest level:
CUDA kernels
CUDA (Compute Unified Device Architecture) is NVIDIA's parallel computing
platform - the lowest-level programming interface for NVIDIA GPUs. When a
framework runs an operation like a matrix multiplication (torch.matmul) or a
convolution, it does not execute a generic CPU instruction; it dispatches a
highly optimized, pre-compiled CUDA kernel to the GPU.
![CUDA data-parallel vector add - arrays A, B, C laid out as columns of cells; each row i, i+1, ... i+n is one thread computing C[i] = A[i] + B[i]](/assets/images/cuda-kernels-944f795026581643555eabdab3a2832f.png)
- Hardware-aware - kernels leverage the GPU's Tensor Cores and memory hierarchy (shared memory, registers).
- Coalesced memory access - memory is accessed in a way that aligns with hardware for maximum bandwidth.
- Massive parallelism - a CPU core handles one or two threads at a time; GPUs have hundreds of cores running many threads at once.
- CPU + GPU together - CUDA makes it easy to use multi-core CPUs and many-core GPUs side by side, playing to each processor's strengths.
XLA compilers
XLA (Accelerated Linear Algebra) is a domain-specific compiler from Google. It began focused on TPUs, but its principles apply to other accelerators including GPUs. Unlike eager execution (each operation runs immediately), XLA works on a computation graph: it captures a whole series of operations, then optimizes the entire graph before compiling. It is analogous to a just-in-time (JIT) compiler, but specialized for linear algebra.
| XLA optimization | What it does |
|---|---|
| Operator fusion | Combines many small operations into a single, larger CUDA kernel - drastically cutting kernel launch overhead and intermediate memory transfers. |
| Memory optimization | Analyzes the graph to minimize memory allocations and deallocations. |
| Layout optimization | Arranges data in memory to suit the hardware's access patterns. |
PyTorch/XLA and JAX/XLA
These integrations pair a framework's user-friendly interface with XLA's compiler optimizations, so developers keep their Pythonic workflow while automatically gaining compiler-level tuning - improving throughput (e.g. tokens/second for LLMs) and reducing cost.
| Integration | How it works |
|---|---|
| PyTorch/XLA | Define the model in standard PyTorch; torch_xla transparently converts the computation graph to XLA's intermediate representation for compilation. Adds gradient checkpointing for memory efficiency and enhanced distributed training (e.g. SPMD support) for large-scale models. |
| JAX/XLA | JAX is built from the ground up with XLA as its core compilation backend. Its functional, tracing-based design is especially well-suited to building and optimizing computation graphs for XLA. |
Achieving performance across diverse hardware
The challenge with mixed environments (GPU, TPU, CPU) is portability: running efficiently on each target without rewriting the code for every device. The same three layers solve this - each one hardware-agnostic where the previous was hardware-specific.
| Layer | Cross-platform role |
|---|---|
| ML frameworks (PyTorch, JAX) | Hardware abstraction. High-level APIs for tensors, operations, and network architectures. Write model code once, then deploy it across accelerators with minimal changes - the framework translates the high-level description into device-specific instructions. |
| XLA compiler | The universal optimizer. XLA is hardware-agnostic: it takes one computation graph and, per target device, uses a specific backend with optimized compilation and code-generation rules. The same graph yields tuned executables for an NVIDIA GPU, a Cloud TPU, or a CPU. |
| PyTorch/XLA & JAX/XLA | The cross-platform enablers. The practical bridge from a unified programming model to different hardware backends. In PyTorch/XLA, device='xla' runs the same code on a TPU, and with SPMD support potentially on a GPU - cutting the engineering overhead of hardware-specific optimization. |
XLA's speedups (fusion, memory and layout optimization) come from analyzing the whole computation graph ahead of time. That is the opposite of eager execution, where each op runs immediately in isolation - so XLA needs the operations captured as a graph before it can compile them.
Recap
The three technologies stack from framework down to hardware, each layer widening portability:
Frameworks like JAX and PyTorch optimize xPU utilization through automatic device placement, efficient memory management, and low-level libraries (CUDA kernels). The XLA compiler adds graph-level optimization that is portable across GPU, TPU, and CPU, while the PyTorch/XLA and JAX/XLA integrations make that portability easy to reach - enabling distributed training (data and model parallelism), memory optimizations, and mixed-precision training for scaling large models.