diff --git a/doc/changelog.rst b/doc/changelog.rst
index 4c82b923..be0dc814 100644
--- a/doc/changelog.rst
+++ b/doc/changelog.rst
@@ -36,6 +36,13 @@ Bug fixes
 - Fixed a bug where the convex hull of a collision mesh was not being computed if the mesh could only collide via a
   :ref:`contact pair<contact-pair>`.
 
+MJX
+^^^
+- Add Warp as a backend implementation for MJX. The implementation can be specified via
+  ``mjx.put_model(m, impl='warp')`` and ``mjx.make_data(m, impl='warp')``. The warp implementation requires
+  a CUDA device and ``warp-lang`` to be installed (``pip install mujoco-mjx[warp]``). This feature is available in
+  "beta" and some bugs are expected.
+
 Version 3.3.4 (July 8, 2025)
 ----------------------------
 
@@ -69,13 +76,6 @@ Documentation
 8. Added missing item documentation and clarified the nature of breaking changes in the 3.3.3 changelog.
    See items 3 and 4 below.
 
-MJX
-^^^
-- Add Warp as a backend implementation for MJX. The implementation can be specified via
-  ``mjx.put_model(m, impl='warp')`` and ``mjx.make_data(m, impl='warp')``. The warp implementation requires
-  a CUDA device, Python 3.12, and `warp-lang` to be installed. This feature is available in "beta" and
-  some bugs are expected.
-
 Version 3.3.3 (June 10, 2025)
 -----------------------------
 
diff --git a/mjx/pyproject.toml b/mjx/pyproject.toml
index b8aeaf44..8bec695f 100644
--- a/mjx/pyproject.toml
+++ b/mjx/pyproject.toml
@@ -33,8 +33,11 @@ dependencies = [
     "mujoco>=3.3.5.dev0",
     "scipy",
     "trimesh",
-    # "warp-lang==1.8.0; python_version >= '3.12' and sys_platform == 'linux'",
-    # "warp-lang==1.8.0; python_version >= '3.12' and sys_platform == 'win32'",
+]
+
+[project.optional-dependencies]
+warp = [
+    "warp-lang==1.8.0",
 ]
 
 [project.scripts]
diff --git a/mjx/requirements.txt b/mjx/requirements.txt
index 9340b6e2..20af0583 100644
--- a/mjx/requirements.txt
+++ b/mjx/requirements.txt
@@ -103,12 +103,7 @@ trimesh==4.5.2 \
     --hash=sha256:2e50f3a7fd135c3045da887a1b9f91230528f3ce11d2ec1ba44750d82d6b4f73
 wheel==0.45.0 \
     --hash=sha256:52f0baa5e6522155090a09c6bd95718cc46956d1b51d537ea5454249edb671c7
-warp-lang==1.8.0; python_version >= '3.12' \
-    --hash=sha256:75a88d2795596f06fcf79eead94e2f194a6195dadbafbd4c9b7c8b4b05456cc4 \
-    --hash=sha256:1be62e7b3e8019ccccc00916c2798afc6c4bfce17d1d6025dc8ef5c790e42c1d \
-    --hash=sha256:373464bee59be37018d134b5924bf8fdb33d1313f484fd1b7c64296971c26ae2 \
-    --hash=sha256:0ecf3b07c1d6d16592ab0318c86f452e16fb3337cfe327bcb76d793a98226dd8 \
-    --hash=sha256:2c7627bee127b522551e02f44c35d7d6a4e5632ff4b58199208012e7f2f6c608
+# TODO(robotics-simulation): add back warp-lang dependencies
 
 # Transitive dependencies of etils[epath]
 fsspec==2024.10.0 \