Improve MJX performance on NVIDIA GPU by enabling Triton GEMM.

PiperOrigin-RevId: 604691509 Change-Id: Iba11a9e1b79848624a1ef6eb4e61e2093d36e662
2024-02-06 10:29:29 -08:00
parent 84f20ffb66
commit 38cbc9e543
3 changed files with 20 additions and 0 deletions
@@ -357,3 +357,13 @@ For MJX to perform well, some configuration parameters should be adjusted from t
  greater, or if MJX detects a TPU as the default backend, otherwise "dense". For TPU, using "sparse" with the
  Newton solver can speed up simulation by 2x to 3x. For GPU, choosing "dense" may impart a more modest speedup of 10%
  to 20%, as long as the dense matrices can fit on the device.
+
+GPU performance tuning
+----------------------
+
+The following environment variables should be set:
+
+``XLA_FLAGS=--xla_gpu_triton_gemm_any=true``
+  This enables the Triton-based GEMM (matmul) emitter for any GEMM that it supports.  This can yield a 30% speedup on
+  NVIDIA GPUs.  If you have multiple GPUs, you may also benefit from enabling flags related to
+  `communciation between GPUs <https://jax.readthedocs.io/en/latest/gpu_performance_tips.html>`__.