Improve MJX performance on NVIDIA GPU by enabling Triton GEMM.

PiperOrigin-RevId: 604691509
Change-Id: Iba11a9e1b79848624a1ef6eb4e61e2093d36e662

mjx/mujoco/mjx/testspeed.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 """Run benchmarks on various devices."""
 
+import os
 import time
 from typing import Sequence, Tuple
 
@@ -56,6 +57,10 @@ def _measure(fn, *args) -> Tuple[float, float]:
 def _main(argv: Sequence[str]):
   """Benchmark a model."""
 
+  xla_flags = os.environ.get('XLA_FLAGS', '')
+  xla_flags += ' --xla_gpu_triton_gemm_any=True'
+  os.environ['XLA_FLAGS'] = xla_flags
+
   f = epath.resource_path('mujoco.mjx') / 'test_data' / FLAGS.mjcf
   m = mujoco.MjModel.from_xml_path(f.as_posix())
   m.opt.solver = {

mjx/tutorial.ipynb

@@ -102,6 +102,11 @@
         "}\n",
         "\"\"\")\n",
         "\n",
+        "# Tell XLA to use Triton GEMM, this improves steps/sec by ~30% on some GPUs\n",
+        "xla_flags = os.environ.get('XLA_FLAGS', '')\n",
+        "xla_flags += ' --xla_gpu_triton_gemm_any=True'\n",
+        "os.environ['XLA_FLAGS'] = xla_flags\n",
+        "\n",
         "# Configure MuJoCo to use the EGL rendering backend (requires GPU)\n",
         "print('Setting environment variable to use GPU rendering:')\n",
         "%env MUJOCO_GL=egl\n",