GPU and CPU performance¶

How probjax's kernels and models compare with the alternatives in JAX and PyTorch, in time and memory, on a consumer NVIDIA GPU and on CPU:

• Attention: dot_product_attention, flex_attention and mha against jax.nn (XLA, cuDNN) and PyTorch SDPA; memory, short sequences, GQA and split-KV decoding, CPU.
• Sequence mixers, kernel sums and MoE: SSD, Mamba scan, residual linear attention, Gated DeltaNet, rbf_kernel_mv / KDE, kernel_logsumexp / Sinkhorn / OT, MoE grouped matmul.
• Models: Transformer / DiffusionTransformer per attention_fn, scan_layers and remat_policy, SSMModel cells, UNet, MoELayer.
• H100 (Hopper): the same suites on an H100, H100 vs RTX 3090, the device-specific attention defaults for sm_90 (cuDNN routing, the Mosaic GPU backend, backward choice), MoE dispatch on H100, and PyTorch (flex_attention, SDPA) on H100.

Devices. These pages measure the RTX 3090 unless they say otherwise; the H100 results are on their own page. Defaults can differ per device: on sm_90 flex_attention / dot_product_attention route large bf16 calls to cuDNN and resolve backward_pass_impl="auto" differently (rules).

At a glance¶

Time of the fastest other implementation that ran (another probjax implementation or an external library) divided by the time of the probjax default, RTX 3090; the CPU row ran on the shared host. Right of the line: the default is faster.

Exact shapes are on the linked pages; ratios at other shapes differ, sometimes a lot.

Setup¶

Every JSON file records the device, the library versions and the runner configuration under device and config.

Methodology¶

• Timing. Each case is compiled once, then timed as inner back-to-back calls bracketed by a device synchronization, repeated several times. Pages show the median per call; every JSON row also has the inter-quartile range (q1_ms, q3_ms). Compilation is timed separately (compile_s) and never included in the step time. No two benchmarks shared the GPU.
• Memory, two kinds. xla_temp_mb / xla_arg_mb / xla_out_mb are XLA's compile-time memory_analysis() of the jitted function: temporaries, arguments, outputs. peak_mb is the true device peak (peak_bytes_in_use) read in a fresh subprocess that runs only that case. Temporaries say what a kernel allocates; the peak says whether the whole call fits. residual_mb is the size of the jax.vjp residuals kept for the backward pass.
• Accuracy is reported as a relative error (max|a − b| / max|b| or a norm ratio; the JSON says which) against an fp32 reference at HIGHEST matmul precision or an fp64 host reference.
• TF32. On Ampere GPUs JAX runs fp32 matmuls in TF32 by default, which costs about three decimal digits; fp32 kernels not pinned to HIGHEST show relative errors around 1e-3.
• Shared host. CPU timings ran on a shared 24-thread machine and carry 10–30% run-to-run noise; read CPU ratios as approximate.
• Cases that did not run. OOM: out of device memory. fail†: XLA found no kernel configuration that fits during autotuning (NOT_FOUND), so the case did not run. unsup.: the implementation rejects the shape or backend. error: another failure (the JSON row has the message). —: not measured. Figures mark them with a cross and the reason; the JSON keeps the error text.
• Figures are generated by docs/scripts/benchmark_figures.py from the JSON, together with these four pages.

What runs by default¶

With no implementation argument, the public functions dispatch as follows. "Kernel" means a Pallas kernel compiled through Triton.

Known limitations¶

• Attention: in some bf16 configurations cuDNN or torch flash is faster than the probjax kernel (bf16 head_dim 128 fwd+bwd at T=8192: probjax 70.9 ms, cuDNN 60.4 ms); the full list is on the attention page.
• Attention on CPU: torch's native SDPA kernel is faster than any XLA formulation (causal fwd+bwd T=4096: probjax 392 ms, torch 101 ms).
• KeyPaddingMask gives wrong results in the attention kernel path; use KVLenMask.
• The RAttention layer runs out of memory at 8192 tokens on a 24 GB GPU because its sliding-window attention branch is dense (kernel page).
• fit() recompiles on every call (models page).
• Dense attention (jax.nn XLA, torch math, probjax implementation='xla') does not run fwd+bwd beyond a few thousand tokens at batch 4 and 8 heads (jax.nn XLA fails from T=8192, torch math from T=4096).

How to reproduce¶

All runners live in docs/scripts/ and write JSON to docs/benchmarks/. GPU runs need a CUDA GPU; torch is optional and only adds comparison rows. Run one benchmark at a time on an otherwise idle GPU, then rebuild the figures and pages.

# Attention grid against JAX and PyTorch, and the tile sweep.
python -m docs.scripts.benchmark_gpu_attention --output docs/benchmarks/gpu_attention.json
python -m docs.scripts.benchmark_gpu_attention_tiles --output docs/benchmarks/gpu_attention_tiles.json

# Model grid.
python -m docs.scripts.benchmark_gpu_nn --output docs/benchmarks/gpu_nn.json

# One kernel family, CPU then GPU (other families: see each page).
JAX_PLATFORMS=cpu python -m docs.scripts.benchmark_kernel_ssd --device cpu \
  --output docs/benchmarks/kernel_ssd_cpu.json
python -m docs.scripts.benchmark_kernel_ssd --device gpu --output docs/benchmarks/kernel_ssd_gpu.json

# Figures (docs/benchmarks/figures/*.svg) and the four GPU pages; needs matplotlib.
python -m docs.scripts.benchmark_figures

Shared timing helpers are in docs/scripts/_gpu_bench.py. Each page ends with the commands for its own files.