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
@@ -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",