GPU performance

GPU performance#

normix runs on a GPU whenever the JAX arrays live there: choose the device with jax.default_device or by placing arrays on it. backend='jax' always runs on the JAX device and backend='cpu' always runs on the CPU; normix never switches between them on its own.

Why small GPU programs are launch-bound#

The Bessel quadrature (log_kv and its derivatives) finds an integration window with fixed-length loops: a few doubling steps and tens of bisection steps. Each step is a tiny GPU kernel, and a GPU launches kernels one after another. The arithmetic is negligible, so the run time is the cost of launching a few hundred kernels in sequence. A CPU has no such cost for a loop of scalar operations, which is why the same code can be faster on the CPU for small inputs.

Optional: XLA command buffers#

XLA can record a sequence of kernels once as a CUDA graph and replay it with a single launch. Loops (WHILE) and branches (CONDITIONAL) are not recorded by default. Enable them before JAX starts:

export XLA_FLAGS="--xla_gpu_enable_command_buffer=+WHILE,+CONDITIONAL --xla_gpu_graph_min_graph_size=1"

or, in Python, before the first JAX operation (before import jax is safest):

import os
os.environ["XLA_FLAGS"] = (
    "--xla_gpu_enable_command_buffer=+WHILE,+CONDITIONAL "
    "--xla_gpu_graph_min_graph_size=1"
)

What the flags mean:

  • --xla_gpu_enable_command_buffer is a comma-separated list of command types that XLA records into CUDA graphs. Setting it replaces the default list (FUSION, CUBLAS, CUBLASLT, CUSTOM_CALL, CUDNN, DYNAMIC_SLICE_FUSION on jaxlib 0.9.1). A + prefix adds a type to the default list instead. An empty value turns command buffers off.

  • --xla_gpu_graph_min_graph_size is the smallest number of kernels worth recording. The default is 5; with 1, short sequences are recorded too.

Measured effect (one RTX 4090 under WSL2, jax 0.9.1, steady-state time of jax.jacrev of GIG.from_expectation; compile time is about 2.2 s in every row):

Setting

Steady-state time

Default

0.0128 s

+WHILE,+CONDITIONAL

0.0100 s

+WHILE,+CONDITIONAL with min_graph_size=1

0.0088 s

Command buffers off

0.0142 s

Results are identical to the printed digits. Batch EM fits of NIG and GH on 5,000 points did not get faster (their time is not dominated by Bessel launch cost) and gave the same final log-likelihood.

Caveats:

  • The flag is global: it applies to every jitted program in the process, not only to the Bessel code.

  • CUDA graphs use extra device memory (about 20 MiB more in the EM test above). If you run out of memory while instantiating graphs, clear the flag.

  • Only one workload on one GPU model was measured. Treat the flag as experimental, and measure your own workload before relying on it.

  • It must be set before JAX initializes the GPU backend; changing it later has no effect.