diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py
index dc10af8c..70d7f973 100644
--- a/mjx/mujoco/mjx/_src/io.py
+++ b/mjx/mujoco/mjx/_src/io.py
@@ -272,6 +272,10 @@ def _put_option(
 
 
   if impl == types.Impl.WARP:
+    if not mjxw.mjwp_io.ENABLE_ISLANDS:
+      fields['disableflags'] = types.DisableBit(
+          fields['disableflags'] | mjwp_types.DisableBit.ISLAND
+      )
     impl_fields = {k: _wp_to_np_type(v) for k, v in impl_fields.items()}
     return types.Option(**fields, _impl=mjxw.types.OptionWarp(**impl_fields))
 
@@ -1272,8 +1276,11 @@ def _get_data_into_warp(
       else:
         value = getattr(d_i, field.name)
 
-      if field.name in ('ne', 'nl', 'nf'):
-        pass
+      if field.name in ('ne', 'nl', 'nf', 'nisland', 'nidof'):
+        if isinstance(value, np.ndarray) and value.size == 0:
+          value = 0
+        else:
+          value = int(value)
       elif field.name in ('nefc', 'ncon'):
         value = {'nefc': nefc, 'ncon': ncon}[field.name]
       elif field.name.endswith('xmat') or field.name == 'ximat':
@@ -1287,9 +1294,11 @@ def _get_data_into_warp(
           'contact',
           'qM',
           'qLD',
+          'qLU',
           'qLDiagInv',
           'ten_J',
           'flexedge_J',
+          'M',
       ):
         continue
       if field.name.startswith('efc_'):
diff --git a/mjx/mujoco/mjx/_src/io_test.py b/mjx/mujoco/mjx/_src/io_test.py
index 21e1910c..bd39f5f4 100644
--- a/mjx/mujoco/mjx/_src/io_test.py
+++ b/mjx/mujoco/mjx/_src/io_test.py
@@ -530,7 +530,7 @@ class DataIOTest(parameterized.TestCase):
     elif impl == 'warp':
       qm = np.zeros((m.nv, m.nv), dtype=np.float64)
       mujoco.mju_sym2dense(qm, d.M, m.M_rownnz, m.M_rowadr, m.M_colind)
-      np.testing.assert_allclose(dx._impl.qM, qm)
+      np.testing.assert_allclose(dx._impl.M[:m.nv, :m.nv], qm)
       # TODO(taylorhowell): test efc__J
       np.testing.assert_allclose(dx._impl.efc__aref[:3], d.efc_aref[:3])
 
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py b/mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py
index 44224545..ec62bc20 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py
@@ -47,6 +47,10 @@ from mujoco.mjx.third_party.mujoco_warp._src.forward import implicit as implicit
 from mujoco.mjx.third_party.mujoco_warp._src.forward import rungekutta4 as rungekutta4
 from mujoco.mjx.third_party.mujoco_warp._src.forward import step1 as step1
 from mujoco.mjx.third_party.mujoco_warp._src.forward import step2 as step2
+from mujoco.mjx.third_party.mujoco_warp._src.history import init_ctrl_history as init_ctrl_history
+from mujoco.mjx.third_party.mujoco_warp._src.history import init_sensor_history as init_sensor_history
+from mujoco.mjx.third_party.mujoco_warp._src.history import read_ctrl as read_ctrl
+from mujoco.mjx.third_party.mujoco_warp._src.history import read_sensor as read_sensor
 from mujoco.mjx.third_party.mujoco_warp._src.inverse import inverse as inverse
 from mujoco.mjx.third_party.mujoco_warp._src.io import create_render_context as create_render_context
 from mujoco.mjx.third_party.mujoco_warp._src.io import get_data_into as get_data_into
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/block_cholesky.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/block_cholesky.py
index 855f9686..7faaa578 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/block_cholesky.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/block_cholesky.py
@@ -19,48 +19,77 @@ import warp as wp
 
 
 @lru_cache(maxsize=None)
-def create_blocked_cholesky_func(block_size: int):
+def create_blocked_cholesky_factorize_solve_func(block_size: int, matrix_size_static: int):
   @wp.func
-  def blocked_cholesky_func(
+  def blocked_cholesky_factorize_solve_func(
     # In:
     A: wp.array2d[float],
+    b: wp.array2d[float],
     matrix_size: int,
     # Out:
-    L: wp.array2d[float],
+    U: wp.array2d[float],
+    x: wp.array2d[float],
   ):
-    """Computes the Cholesky factorization of a symmetric positive definite matrix A in blocks.
+    """Block Cholesky factorization and solve while keeping the forward RHS live."""
+    rhs_tile = wp.tile_load(b, shape=(matrix_size_static, 1), offset=(0, 0), storage="shared", bounds_check=False)
 
-    It returns a lower-triangular matrix L such that A = L L^T.
-    """
-    # Process the matrix in blocks along its leading dimension.
     for k in range(0, matrix_size, block_size):
       end = k + block_size
+      rhs_view = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(k, 0))
 
-      # Load current diagonal block A[k:end, k:end]
-      # and update with contributions from previously computed blocks.
-      A_kk_tile = wp.tile_load(A, shape=(block_size, block_size), offset=(k, k), storage="shared")
+      A_kk_tile = wp.tile_load(
+        A, shape=(block_size, block_size), offset=(k, k), storage="shared", bounds_check=False, aligned=True
+      )
 
       for j in range(0, k, block_size):
-        L_block = wp.tile_load(L, shape=(block_size, block_size), offset=(k, j), storage="shared")
-        wp.tile_matmul(L_block, wp.tile_transpose(L_block), A_kk_tile, alpha=-1.0)
+        U_block = wp.tile_load(
+          U, shape=(block_size, block_size), offset=(j, k), storage="shared", bounds_check=False, aligned=True
+        )
+        wp.tile_matmul(wp.tile_transpose(U_block), U_block, A_kk_tile, alpha=-1.0)
 
-      # Compute the Cholesky factorization for the block
-      wp.tile_cholesky_inplace(A_kk_tile)
-      wp.tile_store(L, A_kk_tile, offset=(k, k))
+        y_block = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(j, 0))
+        wp.tile_matmul(wp.tile_transpose(U_block), y_block, rhs_view, alpha=-1.0)
+
+      wp.tile_cholesky_inplace(A_kk_tile, fill_mode="upper")
+      wp.tile_store(U, A_kk_tile, offset=(k, k), bounds_check=False, aligned=True)
+
+      wp.tile_lower_solve_inplace(wp.tile_transpose(A_kk_tile), rhs_view)
 
-      # Process the blocks below the current block
       for i in range(end, matrix_size, block_size):
-        A_ik_tile = wp.tile_load(A, shape=(block_size, block_size), offset=(i, k), storage="shared")
+        A_ki_tile = wp.tile_load(
+          A, shape=(block_size, block_size), offset=(k, i), storage="shared", bounds_check=False, aligned=True
+        )
 
         for j in range(0, k, block_size):
-          L_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(i, j), storage="shared")
-          L_2_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(k, j), storage="shared")
-          wp.tile_matmul(L_tile, wp.tile_transpose(L_2_tile), A_ik_tile, alpha=-1.0)
+          U_jk_tile = wp.tile_load(
+            U, shape=(block_size, block_size), offset=(j, k), storage="shared", bounds_check=False, aligned=True
+          )
+          U_ji_tile = wp.tile_load(
+            U, shape=(block_size, block_size), offset=(j, i), storage="shared", bounds_check=False, aligned=True
+          )
+          wp.tile_matmul(wp.tile_transpose(U_jk_tile), U_ji_tile, A_ki_tile, alpha=-1.0)
 
-        wp.tile_lower_solve_inplace(A_kk_tile, wp.tile_transpose(A_ik_tile))
-        wp.tile_store(L, A_ik_tile, offset=(i, k))
+        wp.tile_lower_solve_inplace(wp.tile_transpose(A_kk_tile), A_ki_tile)
+        wp.tile_store(U, A_ki_tile, offset=(k, i), bounds_check=False, aligned=True)
 
-  return blocked_cholesky_func
+    for i in range(matrix_size - block_size, -1, -block_size):
+      i_end = i + block_size
+      tmp_tile = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(i, 0))
+      for j in range(i_end, matrix_size, block_size):
+        U_tile = wp.tile_load(
+          U, shape=(block_size, block_size), offset=(i, j), storage="shared", bounds_check=False, aligned=True
+        )
+        x_tile = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(j, 0))
+        wp.tile_matmul(U_tile, x_tile, tmp_tile, alpha=-1.0)
+
+      U_tile = wp.tile_load(
+        U, shape=(block_size, block_size), offset=(i, i), storage="shared", bounds_check=False, aligned=True
+      )
+      wp.tile_upper_solve_inplace(U_tile, tmp_tile)
+
+    wp.tile_store(x, rhs_tile, offset=(0, 0), bounds_check=False)
+
+  return blocked_cholesky_factorize_solve_func
 
 
 @lru_cache(maxsize=None)
@@ -68,41 +97,49 @@ def create_blocked_cholesky_solve_func(block_size: int, matrix_size_static: int)
   @wp.func
   def blocked_cholesky_solve_func(
     # In:
-    L: wp.array2d[float],
+    U: wp.array2d[float],
     b: wp.array2d[float],
     matrix_size: int,
     # Out:
     x: wp.array2d[float],
   ):
-    """Block Cholesky factorization and solve.
+    """Block Cholesky solve.
 
-    Solves A x = b given the Cholesky factor L (A = L L^T) using blocked forward and backward
+    Solves A x = b given the Cholesky factor U (A = U^T U) using blocked forward and backward
     substitution.
     """
     rhs_tile = wp.tile_load(b, shape=(matrix_size_static, 1), offset=(0, 0), storage="shared", bounds_check=False)
 
-    # Forward substitution: solve L y = b
+    # Forward substitution: solve U^T y = b
     for i in range(0, matrix_size, block_size):
       rhs_view = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(i, 0))
       for j in range(0, i, block_size):
-        L_block = wp.tile_load(L, shape=(block_size, block_size), offset=(i, j), storage="shared")
+        U_block = wp.tile_load(
+          U, shape=(block_size, block_size), offset=(j, i), storage="shared", bounds_check=False, aligned=True
+        )
         y_block = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(j, 0))
-        wp.tile_matmul(L_block, y_block, rhs_view, alpha=-1.0)
+        wp.tile_matmul(wp.tile_transpose(U_block), y_block, rhs_view, alpha=-1.0)
 
-      L_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(i, i), storage="shared")
-      wp.tile_lower_solve_inplace(L_tile, rhs_view)
+      U_tile = wp.tile_load(
+        U, shape=(block_size, block_size), offset=(i, i), storage="shared", bounds_check=False, aligned=True
+      )
+      wp.tile_lower_solve_inplace(wp.tile_transpose(U_tile), rhs_view)
 
-    # Backward substitution: solve L^T x = y
+    # Backward substitution: solve U x = y
     for i in range(matrix_size - block_size, -1, -block_size):
       i_end = i + block_size
       tmp_tile = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(i, 0))
       for j in range(i_end, matrix_size, block_size):
-        L_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(j, i), storage="shared")
+        U_tile = wp.tile_load(
+          U, shape=(block_size, block_size), offset=(i, j), storage="shared", bounds_check=False, aligned=True
+        )
         x_tile = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(j, 0))
-        wp.tile_matmul(wp.tile_transpose(L_tile), x_tile, tmp_tile, alpha=-1.0)
-      L_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(i, i), storage="shared")
+        wp.tile_matmul(U_tile, x_tile, tmp_tile, alpha=-1.0)
+      U_tile = wp.tile_load(
+        U, shape=(block_size, block_size), offset=(i, i), storage="shared", bounds_check=False, aligned=True
+      )
 
-      wp.tile_upper_solve_inplace(wp.tile_transpose(L_tile), tmp_tile)
+      wp.tile_upper_solve_inplace(U_tile, tmp_tile)
 
     wp.tile_store(x, rhs_tile, offset=(0, 0), bounds_check=False)
 
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/cli.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/cli.py
index 0e427408..425b54bb 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/cli.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/cli.py
@@ -25,6 +25,7 @@ from absl import flags
 from etils import epath
 
 import mujoco.mjx.third_party.mujoco_warp as mjw
+from mujoco.mjx.third_party.mujoco_warp._src import io
 from mujoco.mjx.third_party.mujoco_warp._src import warp_util
 from mujoco.mjx.third_party.mujoco_warp._src.io import load_trajectory
 from mujoco.mjx.third_party.mujoco_warp._src.io import override_model
@@ -42,6 +43,12 @@ KEYFRAME = flags.DEFINE_integer("keyframe", 0, "keyframe to initialize simulatio
 EVENT_TRACE = flags.DEFINE_bool("event_trace", False, "print an event trace report")
 NOISE_STD = flags.DEFINE_float("noise_std", 0.01, "add noise to ctrl signal (standard deviation)")
 NOISE_RATE = flags.DEFINE_float("noise_rate", 0.1, "add noise to ctrl signal (noise rate)")
+ENABLE_ISLANDS = flags.DEFINE_bool(
+  "enable_islands",
+  False,
+  "Enable constraint islands solver",
+)
+
 
 DEVICE = flags.DEFINE_string("device", None, "override the default Warp device")
 REPLAY = flags.DEFINE_string("replay", None, "NPZ file with ctrl sequence to replay")
@@ -52,6 +59,12 @@ RENDER_RGB = flags.DEFINE_bool("render_rgb", True, "render RGB image")
 RENDER_DEPTH = flags.DEFINE_bool("render_depth", True, "render depth image")
 RENDER_TEXTURES = flags.DEFINE_bool("render_textures", True, "use textures")
 RENDER_SHADOWS = flags.DEFINE_bool("render_shadows", False, "use shadows")
+RENDER_BACKFACE_CULLING = flags.DEFINE_bool(
+  "render_backface_culling",
+  True,
+  "enable renderer backface culling (RenderContext.enable_backface_culling)",
+)
+RENDER_SKYBOX = flags.DEFINE_bool("render_skybox", True, "render skybox texture if available")
 
 
 def load_model(path: epath.Path) -> mujoco.MjModel:
@@ -128,6 +141,8 @@ def init_structs(
   fn: Callable[..., None], mjm: mujoco.MjModel
 ) -> Tuple[mjw.Model, mjw.Data, mjw.RenderContext | None, list[np.ndarray] | None]:
   """Initialize device structs."""
+  io.ENABLE_ISLANDS = ENABLE_ISLANDS.value
+
   mjd = mujoco.MjData(mjm)
   ctrls = None
   if REPLAY.value:
@@ -152,12 +167,14 @@ def init_structs(
 
     rc = mjw.create_render_context(
       mjm,
-      NWORLD.value,
-      (RENDER_WIDTH.value, RENDER_HEIGHT.value),
-      RENDER_RGB.value,
-      RENDER_DEPTH.value,
-      RENDER_TEXTURES.value,
-      RENDER_SHADOWS.value,
+      nworld=NWORLD.value,
+      cam_res=(RENDER_WIDTH.value, RENDER_HEIGHT.value),
+      render_rgb=RENDER_RGB.value,
+      render_depth=RENDER_DEPTH.value,
+      use_textures=RENDER_TEXTURES.value,
+      use_shadows=RENDER_SHADOWS.value,
+      enable_backface_culling=RENDER_BACKFACE_CULLING.value,
+      render_skybox=RENDER_SKYBOX.value,
     )
 
     return m, d, rc, ctrls
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py
index a444e399..9bf7ce5e 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py
@@ -23,6 +23,8 @@ from mujoco.mjx.third_party.mujoco_warp._src.collision_core import contact_param
 from mujoco.mjx.third_party.mujoco_warp._src.collision_core import geom_collision_pair
 from mujoco.mjx.third_party.mujoco_warp._src.collision_core import write_contact
 from mujoco.mjx.third_party.mujoco_warp._src.collision_gjk import ccd
+from mujoco.mjx.third_party.mujoco_warp._src.collision_gjk import epa_phase
+from mujoco.mjx.third_party.mujoco_warp._src.collision_gjk import gjk_phase
 from mujoco.mjx.third_party.mujoco_warp._src.collision_gjk import multicontact
 from mujoco.mjx.third_party.mujoco_warp._src.collision_gjk import support
 from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import Geom
@@ -35,7 +37,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAX_EPAFACES
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAX_EPAHORIZON
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXCONPAIR
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
-from mujoco.mjx.third_party.mujoco_warp._src.types import _NEW_GAP_SEMANTICS
+from mujoco.mjx.third_party.mujoco_warp._src.types import NEW_GAP_SEMANTICS
 from mujoco.mjx.third_party.mujoco_warp._src.types import Data
 from mujoco.mjx.third_party.mujoco_warp._src.types import DisableBit
 from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
@@ -715,6 +717,7 @@ def ccd_kernel_builder(
     opt_ccd_tolerance: wp.array[float],
     # Data in:
     naconmax_in: int,
+    naccdmax_in: int,
     # In:
     epa_vert_in: wp.array2d[wp.vec3],
     epa_vert_index_in: wp.array2d[int],
@@ -737,7 +740,7 @@ def ccd_kernel_builder(
     geom2: Geom,
     geoms: wp.vec2i,
     worldid: int,
-    ccdid: int,
+    nccd_in: wp.array[int],
     margin: float,
     gap: float,
     condim: int,
@@ -774,30 +777,47 @@ def ccd_kernel_builder(
     if is_collision_sensor:
       cutoff = 1.0e32
     else:
-      if wp.static(_NEW_GAP_SEMANTICS):
+      if wp.static(NEW_GAP_SEMANTICS):
         cutoff = gap
       else:
         cutoff = 0.0
-    dist, ncollision, w1, w2, multiccd_idx = ccd(
+    needs_epa, dist, ncollision, w1, w2, gjk_result, geom1, geom2 = gjk_phase(
       opt_ccd_tolerance[worldid % opt_ccd_tolerance.shape[0]],
       cutoff,
       gjk_iterations,
-      epa_iterations,
       geom1,
       geom2,
       geomtype1,
       geomtype2,
       x1,
       x2,
-      epa_vert_in[ccdid],
-      epa_vert_index_in[ccdid],
-      epa_face_in[ccdid],
-      epa_pr_in[ccdid],
-      epa_norm2_in[ccdid],
-      epa_horizon_in[ccdid],
     )
 
-    if wp.static(_NEW_GAP_SEMANTICS):
+    ccdid = int(-1)
+    multiccd_idx = int(-1)
+
+    if needs_epa:
+      ccdid = wp.atomic_add(nccd_in, geomgeomid, 1)
+      if ccdid >= naccdmax_in:
+        wp.printf("CCD overflow - please increase naccdmax to %u\n", ccdid)
+        return 0
+      dist, ncollision, w1, w2, multiccd_idx = epa_phase(
+        opt_ccd_tolerance[worldid % opt_ccd_tolerance.shape[0]],
+        epa_iterations,
+        gjk_result,
+        geom1,
+        geom2,
+        geomtype1,
+        geomtype2,
+        epa_vert_in[ccdid],
+        epa_vert_index_in[ccdid],
+        epa_face_in[ccdid],
+        epa_pr_in[ccdid],
+        epa_norm2_in[ccdid],
+        epa_horizon_in[ccdid],
+      )
+
+    if wp.static(NEW_GAP_SEMANTICS):
       if dist >= gap and pairid[1] == -1:
         return 0
     else:
@@ -990,11 +1010,6 @@ def ccd_kernel_builder(
     if geom_type[g1] != geomtype1 or geom_type[g2] != geomtype2:
       return
 
-    ccdid = wp.atomic_add(nccd_in, wp.static(geomgeomid), 1)
-    if ccdid >= naccdmax_in:
-      wp.printf("CCD overflow - please increase naccdmax to %u\n", ccdid)
-      return
-
     worldid = collision_worldid_in[collisionid]
 
     _, margin, gap, condim, friction, solref, solreffriction, solimp = contact_params(
@@ -1046,6 +1061,7 @@ def ccd_kernel_builder(
     eval_ccd_write_contact(
       opt_ccd_tolerance,
       naconmax_in,
+      naccdmax_in,
       epa_vert_in,
       epa_vert_index_in,
       epa_face_in,
@@ -1067,7 +1083,7 @@ def ccd_kernel_builder(
       geom2,
       geoms,
       worldid,
-      ccdid,
+      nccd_in,
       margin,
       gap,
       condim,
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_core.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_core.py
index 41feb836..35bb0f27 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_core.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_core.py
@@ -22,8 +22,8 @@ import warp as wp
 
 from mujoco.mjx.third_party.mujoco_warp._src.math import safe_div
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINMU
-from mujoco.mjx.third_party.mujoco_warp._src.types import _NEW_GAP_SEMANTICS
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
+from mujoco.mjx.third_party.mujoco_warp._src.types import NEW_GAP_SEMANTICS
 from mujoco.mjx.third_party.mujoco_warp._src.types import ContactType
 from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
 from mujoco.mjx.third_party.mujoco_warp._src.types import mat63
@@ -219,7 +219,7 @@ def write_contact(
     contact_frame_out[cid] = frame_in
     contact_geom_out[cid] = geoms_in
     contact_worldid_out[cid] = worldid_in
-    if wp.static(_NEW_GAP_SEMANTICS):
+    if wp.static(NEW_GAP_SEMANTICS):
       includemargin = margin_in
     else:
       includemargin = margin_in - gap_in
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py
index df75adc7..7c2fe28e 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py
@@ -375,6 +375,7 @@ def _binary_search(values: wp.array[Any], value: Any, lower: int, upper: int) ->
   return upper
 
 
+@cache_kernel
 def _sap_project(opt_broadphase: int):
   @wp.kernel(module="unique", enable_backward=False)
   def sap_project(
@@ -531,6 +532,7 @@ def _sap_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: i
   return kernel
 
 
+@cache_kernel
 def _segmented_sort(tile_size: int):
   @wp.kernel(module="unique")
   def segmented_sort(
@@ -629,7 +631,9 @@ def sap_broadphase(m: Model, d: Data, ctx: CollisionContext):
   # assumes each geom has 5 other geoms (batched over all worlds)
   nsweep = 5 * nworldgeom
   wp.launch(
-    kernel=_sap_broadphase(m.opt.broadphase_filter, m.geom_aabb.shape[0], m.geom_rbound.shape[0], m.geom_margin.shape[0], m.geom_gap.shape[0]),
+    kernel=_sap_broadphase(
+      m.opt.broadphase_filter, m.geom_aabb.shape[0], m.geom_rbound.shape[0], m.geom_margin.shape[0], m.geom_gap.shape[0]
+    ),
     dim=nsweep,
     inputs=[
       m.ngeom,
@@ -718,7 +722,9 @@ def nxn_broadphase(m: Model, d: Data, ctx: CollisionContext):
   `contype`/`conaffinity`, parent-child relationships, and explicit `<exclude>` tags.
   """
   wp.launch(
-    _nxn_broadphase(m.opt.broadphase_filter, m.geom_aabb.shape[0], m.geom_rbound.shape[0], m.geom_margin.shape[0], m.geom_gap.shape[0]),
+    _nxn_broadphase(
+      m.opt.broadphase_filter, m.geom_aabb.shape[0], m.geom_rbound.shape[0], m.geom_margin.shape[0], m.geom_gap.shape[0]
+    ),
     dim=(d.nworld, m.nxn_geom_pair_filtered.shape[0]),
     inputs=[
       m.geom_type,
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk.py
index 21e61b7c..f299cd78 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk.py
@@ -2135,8 +2135,12 @@ def multicontact(
 
   # face1 is an edge; clip face1 against face2
   if is_edge_contact_geom1:
-    approx_dir = wp.norm_l2(dir) * n2[j]
-    return _polygon_clip(plane_normal, plane_dist, face2, nface2, face1, nface1, n2[j], approx_dir, polygon, clipped)
+    approx_dir = -wp.norm_l2(dir) * n2[j]
+    nclipped, clipped1, clipped2 = _polygon_clip(
+      plane_normal, plane_dist, face2, nface2, face1, nface1, n2[j], approx_dir, polygon, clipped
+    )
+    # the faces were flipped in calling _polygon_clip so we need to flip them back
+    return nclipped, clipped2, clipped1
 
   # face2 is an edge; clip face2 against face1
   if is_edge_contact_geom2:
@@ -2197,34 +2201,29 @@ def _inflate(
 
 
 @wp.func
-def ccd(
+def gjk_phase(
   # In:
   tolerance: float,
   cutoff: float,
   gjk_iterations: int,
-  epa_iterations: int,
   geom1: Geom,
   geom2: Geom,
   geomtype1: int,
   geomtype2: int,
   x_1: wp.vec3,
   x_2: wp.vec3,
-  vert: wp.array[wp.vec3],
-  vert_index: wp.array[int],
-  face: wp.array[int],
-  face_pr: wp.array[wp.vec3],
-  face_norm2: wp.array[float],
-  horizon: wp.array[int],
-) -> Tuple[float, int, wp.vec3, wp.vec3, int]:
-  """General convex collision detection via GJK/EPA."""
+) -> Tuple[bool, float, int, wp.vec3, wp.vec3, GJKResult, Geom, Geom]:
+  """Run GJK phase of CCD."""
   full_margin1 = 0.0
   full_margin2 = 0.0
   size1 = 0.0
   size2 = 0.0
+  empty = GJKResult()
 
   # determine if the geoms being tested are discrete
   is_discrete = _discrete_geoms(geomtype1, geomtype2) and (geom1.margin == 0.0 and geom2.margin == 0.0)
 
+  # special handling for sphere and capsule (shrink to point and line respectively)
   if geomtype1 == GeomType.SPHERE or geomtype1 == GeomType.CAPSULE:
     size1 = geom1.size[0]
     full_margin1 = size1 + 0.5 * geom1.margin
@@ -2237,7 +2236,6 @@ def ccd(
     geom2.margin = 0.0
     geom2.size = wp.vec3(0.0, geom2.size[1], geom2.size[2])
 
-  # special handling for sphere and capsule (shrink to point and line respectively)
   if size1 + size2 > 0.0:
     cutoff += full_margin1 + full_margin2
     result = gjk(tolerance, gjk_iterations, geom1, geom2, x_1, x_2, geomtype1, geomtype2, cutoff, is_discrete)
@@ -2245,11 +2243,11 @@ def ccd(
     # shallow penetration, inflate contact
     if result.dist > tolerance:
       if result.dist == FLOAT_MAX:
-        return result.dist, 1, result.x1, result.x2, -1
+        return False, result.dist, 1, result.x1, result.x2, empty, geom1, geom2
       dist, x1, x2 = _inflate(result, geom1, geom2, geomtype1, geomtype2, full_margin1, full_margin2)
-      return dist, 1, x1, x2, -1
+      return False, dist, 1, x1, x2, empty, geom1, geom2
 
-    # deep penetration, reset initial conditions and rerun GJK + EPA
+    # deep penetration: reset initial conditions and rerun GJK + EPA
     geom1.margin = full_margin1 - size1
     geom1.size = wp.vec3(size1, geom1.size[1], geom1.size[2])
     geom2.margin = full_margin2 - size2
@@ -2260,8 +2258,29 @@ def ccd(
 
   # no penetration depth to recover
   if result.dist > tolerance or result.dim < 2:
-    return result.dist, 1, result.x1, result.x2, -1
+    return False, result.dist, 1, result.x1, result.x2, empty, geom1, geom2
 
+  return True, result.dist, 1, result.x1, result.x2, result, geom1, geom2
+
+
+@wp.func
+def epa_phase(
+  # In:
+  tolerance: float,
+  epa_iterations: int,
+  result: GJKResult,
+  geom1: Geom,
+  geom2: Geom,
+  geomtype1: int,
+  geomtype2: int,
+  vert: wp.array[wp.vec3],
+  vert_index: wp.array[int],
+  face: wp.array[int],
+  face_pr: wp.array[wp.vec3],
+  face_norm2: wp.array[float],
+  horizon: wp.array[int],
+) -> Tuple[float, int, wp.vec3, wp.vec3, int]:
+  """Run EPA given GJK result. Returns (dist, ncontact, x1, x2, multiccd_idx)."""
   pt = Polytope()
   pt.nface = 0
   pt.nvert = 0
@@ -2330,6 +2349,7 @@ def ccd(
   if pt.status:
     return result.dist, 1, result.x1, result.x2, -1
 
+  is_discrete = _discrete_geoms(geomtype1, geomtype2) and (geom1.margin == 0.0 and geom2.margin == 0.0)
   dist, x1, x2, idx = _epa(tolerance, epa_iterations, pt, geom1, geom2, geomtype1, geomtype2, is_discrete)
   if idx == -1:
     return FLOAT_MAX, 0, wp.vec3(), wp.vec3(), -1
@@ -2343,3 +2363,46 @@ def ccd(
     idx = -1
 
   return dist, 1, x1, x2, idx
+
+
+@wp.func
+def ccd(
+  # In:
+  tolerance: float,
+  cutoff: float,
+  gjk_iterations: int,
+  epa_iterations: int,
+  geom1: Geom,
+  geom2: Geom,
+  geomtype1: int,
+  geomtype2: int,
+  x_1: wp.vec3,
+  x_2: wp.vec3,
+  vert: wp.array[wp.vec3],
+  vert_index: wp.array[int],
+  face: wp.array[int],
+  face_pr: wp.array[wp.vec3],
+  face_norm2: wp.array[float],
+  horizon: wp.array[int],
+) -> Tuple[float, int, wp.vec3, wp.vec3, int]:
+  """General convex collision detection via GJK/EPA."""
+  needs_epa, dist, ncontact, x1, x2, result, geom1, geom2 = gjk_phase(
+    tolerance, cutoff, gjk_iterations, geom1, geom2, geomtype1, geomtype2, x_1, x_2
+  )
+  if not needs_epa:
+    return dist, ncontact, x1, x2, -1
+  return epa_phase(
+    tolerance,
+    epa_iterations,
+    result,
+    geom1,
+    geom2,
+    geomtype1,
+    geomtype2,
+    vert,
+    vert_index,
+    face,
+    face_pr,
+    face_norm2,
+    horizon,
+  )
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_sdf.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_sdf.py
index 5937cf6e..a96c8985 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_sdf.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_sdf.py
@@ -23,6 +23,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.collision_core import geom_collisio
 from mujoco.mjx.third_party.mujoco_warp._src.collision_core import write_contact
 from mujoco.mjx.third_party.mujoco_warp._src.math import make_frame
 from mujoco.mjx.third_party.mujoco_warp._src.ray import ray_mesh
+from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
 from mujoco.mjx.third_party.mujoco_warp._src.types import Data
 from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
 from mujoco.mjx.third_party.mujoco_warp._src.types import Model
@@ -183,6 +184,24 @@ def ellipsoid(p: wp.vec3, size: wp.vec3) -> float:
   return k0 * (k0 - 1.0) / denom
 
 
+@wp.func
+def capsule(p: wp.vec3, size: wp.vec3) -> float:
+  r = size[0]
+  h = size[1]
+  pz_clamped = wp.clamp(p[2], -h, h)
+  diff = wp.vec3(p[0], p[1], p[2] - pz_clamped)
+  return wp.length(diff) - r
+
+
+@wp.func
+def cylinder(p: wp.vec3, size: wp.vec3) -> float:
+  r = size[0]
+  h = size[1]
+  dx = wp.length(wp.vec2(p[0], p[1])) - r
+  dy = wp.abs(p[2]) - h
+  return wp.min(wp.max(dx, dy), 0.0) + wp.length(wp.vec2(wp.max(dx, 0.0), wp.max(dy, 0.0)))
+
+
 @wp.func
 def grad_sphere(p: wp.vec3) -> wp.vec3:
   c = wp.length(p)
@@ -230,6 +249,52 @@ def grad_ellipsoid(p: wp.vec3, size: wp.vec3) -> wp.vec3:
   return raw_grad / wp.length(raw_grad)
 
 
+@wp.func
+def grad_capsule(p: wp.vec3, size: wp.vec3) -> wp.vec3:
+  h = size[1]
+  pz_clamped = wp.clamp(p[2], -h, h)
+  diff = wp.vec3(p[0], p[1], p[2] - pz_clamped)
+  c = wp.length(diff)
+  if c > MJ_MINVAL:
+    return diff / c
+  else:
+    return wp.vec3(0.0)
+
+
+@wp.func
+def grad_cylinder(p: wp.vec3, size: wp.vec3) -> wp.vec3:
+  r = size[0]
+  h = size[1]
+
+  radial_dist = wp.length(wp.vec2(p[0], p[1]))
+  if radial_dist > MJ_MINVAL:
+    u = wp.vec3(p[0] / radial_dist, p[1] / radial_dist, 0.0)
+  else:
+    u = wp.vec3(0.0)
+
+  w = wp.vec3(0.0, 0.0, wp.sign(p[2]))
+
+  dx = radial_dist - r
+  dy = wp.abs(p[2]) - h
+
+  if dx > 0.0 and dy > 0.0:
+    v = wp.vec2(dx, dy)
+    len_v = wp.length(v)
+    if len_v > MJ_MINVAL:
+      return u * (dx / len_v) + w * (dy / len_v)
+    else:
+      return wp.vec3(0.0)
+  elif dx > 0.0:
+    return u
+  elif dy > 0.0:
+    return w
+  else:
+    if dx > dy:
+      return u
+    else:
+      return w
+
+
 @wp.func
 def user_sdf(p: wp.vec3, attr: vec_pluginattr, sdf_type: int) -> float:
   """User-defined SDF function.
@@ -287,16 +352,17 @@ def find_oct(
     # check if the node is a leaf
     # child indices are relative to root (mesh_octadr offset)
     child0 = oct_child[node][0]
+    # Evaluate this hot leaf predicate eagerly to avoid branch-heavy codegen.
     if (
-      child0 == -1
-      and oct_child[node][1] == -1
-      and oct_child[node][2] == -1
-      and oct_child[node][3] == -1
-      and oct_child[node][4] == -1
-      and oct_child[node][5] == -1
-      and oct_child[node][6] == -1
-      and oct_child[node][7] == -1
-    ):
+      int(child0 == -1)
+      & int(oct_child[node][1] == -1)
+      & int(oct_child[node][2] == -1)
+      & int(oct_child[node][3] == -1)
+      & int(oct_child[node][4] == -1)
+      & int(oct_child[node][5] == -1)
+      & int(oct_child[node][6] == -1)
+      & int(oct_child[node][7] == -1)
+    ) != 0:
       for j in range(8):
         if not grad:
           rx[j] = (
@@ -394,6 +460,10 @@ def sdf(type: int, p: wp.vec3, attr: vec_pluginattr, sdf_type: int, volume_data:
     return p[2]
   elif type == GeomType.SPHERE:
     return sphere(p, attr_vec3)
+  elif type == GeomType.CAPSULE:
+    return capsule(p, attr_vec3)
+  elif type == GeomType.CYLINDER:
+    return cylinder(p, attr_vec3)
   elif type == GeomType.BOX:
     return box(p, attr_vec3)
   elif type == GeomType.ELLIPSOID:
@@ -452,6 +522,10 @@ def sdf_grad(
     return grad
   elif type == GeomType.SPHERE:
     return grad_sphere(p)
+  elif type == GeomType.CAPSULE:
+    return grad_capsule(p, attr_vec3)
+  elif type == GeomType.CYLINDER:
+    return grad_cylinder(p, attr_vec3)
   elif type == GeomType.BOX:
     return grad_box(p, attr_vec3)
   elif type == GeomType.ELLIPSOID:
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py
index 59a7ad72..bcb1acf5 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py
@@ -134,7 +134,10 @@ def _equality_connect(
   body_dofnum: wp.array[int],
   body_dofadr: wp.array[int],
   body_invweight0: wp.array2d[wp.vec2],
+  jnt_type: wp.array[int],
+  jnt_dofadr: wp.array[int],
   dof_bodyid: wp.array[int],
+  dof_jntid: wp.array[int],
   dof_parentid: wp.array[int],
   site_bodyid: wp.array[int],
   eq_obj1id: wp.array[int],
@@ -154,6 +157,9 @@ def _equality_connect(
   site_xpos_in: wp.array2d[wp.vec3],
   subtree_com_in: wp.array2d[wp.vec3],
   cdof_in: wp.array2d[wp.spatial_vector],
+  cvel_in: wp.array2d[wp.spatial_vector],
+  cdof_dot_in: wp.array2d[wp.spatial_vector],
+  subtree_linvel_in: wp.array2d[wp.vec3],
   njmax_in: int,
   njmax_nnz_in: int,
   # Data out:
@@ -214,6 +220,7 @@ def _equality_connect(
 
   # compute Jacobian difference (opposite of contact: 0 - 1)
   Jqvel = wp.vec3f(0.0, 0.0, 0.0)
+  Jdotv = wp.vec3f(0.0, 0.0, 0.0)
 
   if is_sparse:
     # TODO(team): pre-compute number of non-zeros
@@ -282,6 +289,42 @@ def _equality_connect(
       )
       j1mj2 = jacp1 - jacp2
 
+      jacp1_dot, _ = support.jac_dot_dof(
+        body_parentid,
+        body_rootid,
+        jnt_type,
+        jnt_dofadr,
+        dof_bodyid,
+        dof_jntid,
+        body_isdofancestor,
+        subtree_com_in,
+        cdof_in,
+        cvel_in,
+        cdof_dot_in,
+        pos1,
+        body1,
+        da,
+        worldid,
+      )
+      jacp2_dot, _ = support.jac_dot_dof(
+        body_parentid,
+        body_rootid,
+        jnt_type,
+        jnt_dofadr,
+        dof_bodyid,
+        dof_jntid,
+        body_isdofancestor,
+        subtree_com_in,
+        cdof_in,
+        cvel_in,
+        cdof_dot_in,
+        pos2,
+        body2,
+        da,
+        worldid,
+      )
+      j1mj2_dot = jacp1_dot - jacp2_dot
+
       sparseid0 = rowadr + nnz
       sparseid1 = rowadr + rownnz + nnz
       sparseid2 = rowadr + 2 * rownnz + nnz
@@ -294,7 +337,9 @@ def _equality_connect(
       efc_J_out[worldid, 0, sparseid1] = j1mj2[1]
       efc_J_out[worldid, 0, sparseid2] = j1mj2[2]
 
-      Jqvel += j1mj2 * qvel_in[worldid, da]
+      qvel = qvel_in[worldid, da]
+      Jqvel += j1mj2 * qvel
+      Jdotv += j1mj2_dot * qvel
 
       nnz += 1
   else:
@@ -326,11 +371,49 @@ def _equality_connect(
       )
       j1mj2 = jacp1 - jacp2
 
+      jacp1_dot, _ = support.jac_dot_dof(
+        body_parentid,
+        body_rootid,
+        jnt_type,
+        jnt_dofadr,
+        dof_bodyid,
+        dof_jntid,
+        body_isdofancestor,
+        subtree_com_in,
+        cdof_in,
+        cvel_in,
+        cdof_dot_in,
+        pos1,
+        body1,
+        dofid,
+        worldid,
+      )
+      jacp2_dot, _ = support.jac_dot_dof(
+        body_parentid,
+        body_rootid,
+        jnt_type,
+        jnt_dofadr,
+        dof_bodyid,
+        dof_jntid,
+        body_isdofancestor,
+        subtree_com_in,
+        cdof_in,
+        cvel_in,
+        cdof_dot_in,
+        pos2,
+        body2,
+        dofid,
+        worldid,
+      )
+      j1mj2_dot = jacp1_dot - jacp2_dot
+
       efc_J_out[worldid, efcid0, dofid] = j1mj2[0]
       efc_J_out[worldid, efcid1, dofid] = j1mj2[1]
       efc_J_out[worldid, efcid2, dofid] = j1mj2[2]
 
-      Jqvel += j1mj2 * qvel_in[worldid, dofid]
+      qvel = qvel_in[worldid, dofid]
+      Jqvel += j1mj2 * qvel
+      Jdotv += j1mj2_dot * qvel
 
   body_invweight0_id = worldid % body_invweight0.shape[0]
   invweight = body_invweight0[body_invweight0_id, body1][0] + body_invweight0[body_invweight0_id, body2][0]
@@ -368,6 +451,8 @@ def _equality_connect(
       efc_frictionloss_out,
     )
 
+    efc_aref_out[worldid, efcidi] -= Jdotv[i]
+
 
 @wp.kernel
 def _equality_joint(
@@ -812,7 +897,10 @@ def _equality_weld(
   body_dofnum: wp.array[int],
   body_dofadr: wp.array[int],
   body_invweight0: wp.array2d[wp.vec2],
+  jnt_type: wp.array[int],
+  jnt_dofadr: wp.array[int],
   dof_bodyid: wp.array[int],
+  dof_jntid: wp.array[int],
   dof_parentid: wp.array[int],
   site_bodyid: wp.array[int],
   site_quat: wp.array2d[wp.quat],
@@ -834,6 +922,9 @@ def _equality_weld(
   site_xpos_in: wp.array2d[wp.vec3],
   subtree_com_in: wp.array2d[wp.vec3],
   cdof_in: wp.array2d[wp.spatial_vector],
+  cvel_in: wp.array2d[wp.spatial_vector],
+  cdof_dot_in: wp.array2d[wp.spatial_vector],
+  subtree_linvel_in: wp.array2d[wp.vec3],
   njmax_in: int,
   njmax_nnz_in: int,
   # Data out:
@@ -903,9 +994,43 @@ def _equality_weld(
     quat = math.mul_quat(xquat_in[worldid, body1], relpose)
     quat1 = math.quat_inv(xquat_in[worldid, body2])
 
+    # quat1 = quat_inv(xquat_in[worldid, body2])
+    q2 = xquat_in[worldid, body2]
+    quat1 = wp.quat(q2[0], -q2[1], -q2[2], -q2[3])
+
+  # compute rotational Jdotv helper quaternions
+  omega1 = wp.spatial_top(cvel_in[worldid, body1])
+  omega2 = wp.spatial_top(cvel_in[worldid, body2])
+  domega = omega1 - omega2
+
+  omega1_q = wp.quat(0.0, omega1[0], omega1[1], omega1[2])
+  omega2_q = wp.quat(0.0, omega2[0], omega2[1], omega2[2])
+  domega_q = wp.quat(0.0, domega[0], domega[1], domega[2])
+
+  if is_site:
+    qdot0r = math.mul_quat(omega1_q, quat) * 0.5
+    qfull1 = math.mul_quat(xquat_in[worldid, body2], site_quat[site_quat_id, obj2id])
+    qdot1 = math.mul_quat(omega2_q, qfull1) * 0.5
+
+    negqdot1 = wp.quat(-qdot1[0], -qdot1[1], -qdot1[2], -qdot1[3])
+    negq1 = wp.quat(qfull1[0], -qfull1[1], -qfull1[2], -qfull1[3])
+
+  else:
+    # qdot0 = mul_quat(xquat_in[worldid, body1], omega1_q) * 0.5
+    u7 = xquat_in[worldid, body1]
+    qdot0 = math.mul_quat(omega1_q, xquat_in[worldid, body1]) * 0.5
+    qdot0r = math.mul_quat(qdot0, relpose)
+    q1_non_site = xquat_in[worldid, body2]
+    qdot1 = math.mul_quat(omega2_q, q1_non_site) * 0.5
+
+    negqdot1 = wp.quat(-qdot1[0], -qdot1[1], -qdot1[2], -qdot1[3])
+    negq1 = wp.quat(q1_non_site[0], -q1_non_site[1], -q1_non_site[2], -q1_non_site[3])
+
   # compute Jacobian difference (opposite of contact: 0 - 1)
   Jqvelp = wp.vec3f(0.0, 0.0, 0.0)
   Jqvelr = wp.vec3f(0.0, 0.0, 0.0)
+  Jdotv_p = wp.vec3f(0.0, 0.0, 0.0)
+  Jdotv_r0 = wp.vec3f(0.0, 0.0, 0.0)
 
   if is_sparse:
     # TODO(team): pre-compute number of non-zeros
@@ -985,6 +1110,44 @@ def _equality_weld(
       jacdifrq = math.mul_quat(math.quat_mul_axis(quat1, jacdifr), quat)
       jacdifr = 0.5 * wp.vec3(jacdifrq[1], jacdifrq[2], jacdifrq[3])
 
+      jacp1_dot, jacr1_dot = support.jac_dot_dof(
+        body_parentid,
+        body_rootid,
+        jnt_type,
+        jnt_dofadr,
+        dof_bodyid,
+        dof_jntid,
+        body_isdofancestor,
+        subtree_com_in,
+        cdof_in,
+        cvel_in,
+        cdof_dot_in,
+        pos1,
+        body1,
+        da,
+        worldid,
+      )
+      jacp2_dot, jacr2_dot = support.jac_dot_dof(
+        body_parentid,
+        body_rootid,
+        jnt_type,
+        jnt_dofadr,
+        dof_bodyid,
+        dof_jntid,
+        body_isdofancestor,
+        subtree_com_in,
+        cdof_in,
+        cvel_in,
+        cdof_dot_in,
+        pos2,
+        body2,
+        da,
+        worldid,
+      )
+
+      jacdifp_dot = jacp1_dot - jacp2_dot
+      jacdifr_dot = jacr1_dot - jacr2_dot
+
       sparseid0 = rowadr + nnz
       sparseid1 = rowadr + rownnz + nnz
       sparseid2 = rowadr + 2 * rownnz + nnz
@@ -1008,6 +1171,8 @@ def _equality_weld(
 
       Jqvelp += jacdifp * qvel_in[worldid, da]
       Jqvelr += jacdifr * qvel_in[worldid, da]
+      Jdotv_p += jacdifp_dot * qvel_in[worldid, da]
+      Jdotv_r0 += jacdifr_dot * qvel_in[worldid, da]
 
       nnz += 1
   else:
@@ -1047,12 +1212,52 @@ def _equality_weld(
       jacdifrq = math.mul_quat(math.quat_mul_axis(quat1, jacdifr), quat)
       jacdifr = 0.5 * wp.vec3(jacdifrq[1], jacdifrq[2], jacdifrq[3])
 
+      jacp1_dot, jacr1_dot = support.jac_dot_dof(
+        body_parentid,
+        body_rootid,
+        jnt_type,
+        jnt_dofadr,
+        dof_bodyid,
+        dof_jntid,
+        body_isdofancestor,
+        subtree_com_in,
+        cdof_in,
+        cvel_in,
+        cdof_dot_in,
+        pos1,
+        body1,
+        dofid,
+        worldid,
+      )
+      jacp2_dot, jacr2_dot = support.jac_dot_dof(
+        body_parentid,
+        body_rootid,
+        jnt_type,
+        jnt_dofadr,
+        dof_bodyid,
+        dof_jntid,
+        body_isdofancestor,
+        subtree_com_in,
+        cdof_in,
+        cvel_in,
+        cdof_dot_in,
+        pos2,
+        body2,
+        dofid,
+        worldid,
+      )
+
+      jacdifp_dot = jacp1_dot - jacp2_dot
+      jacdifr_dot = jacr1_dot - jacr2_dot
+
       efc_J_out[worldid, efcid3, dofid] = jacdifr[0]
       efc_J_out[worldid, efcid4, dofid] = jacdifr[1]
       efc_J_out[worldid, efcid5, dofid] = jacdifr[2]
 
       Jqvelp += jacdifp * qvel_in[worldid, dofid]
       Jqvelr += jacdifr * qvel_in[worldid, dofid]
+      Jdotv_p += jacdifp_dot * qvel_in[worldid, dofid]
+      Jdotv_r0 += jacdifr_dot * qvel_in[worldid, dofid]
 
   # error is difference in global position and orientation
   cpos = pos1 - pos2
@@ -1070,6 +1275,22 @@ def _equality_weld(
 
   timestep = opt_timestep[worldid % opt_timestep.shape[0]]
 
+  djrdv_q = wp.quat(0.0, Jdotv_r0[0], Jdotv_r0[1], Jdotv_r0[2])
+
+  # Term 1: negqdot1 * domega * q0r
+  t1a = math.mul_quat(negqdot1, domega_q)
+  t1 = math.mul_quat(t1a, quat)
+
+  # Term 2: negq1 * djrdv * q0r
+  t2a = math.mul_quat(negq1, djrdv_q)
+  t2 = math.mul_quat(t2a, quat)
+
+  # Term 3: negq1 * domega * qdot0r
+  t3a = math.mul_quat(negq1, domega_q)
+  t3 = math.mul_quat(t3a, qdot0r)
+
+  Jdotv_r = wp.vec3(t1[1] + t2[1] + t3[1], t1[2] + t2[2] + t3[2], t1[3] + t2[3] + t3[3]) * 0.5 * torquescale
+
   for i in range(3):
     _efc_row(
       opt_disableflags,
@@ -1096,6 +1317,8 @@ def _equality_weld(
       efc_frictionloss_out,
     )
 
+    efc_aref_out[worldid, efcid + i] -= Jdotv_p[i]
+
   invweight_r = body_invweight0[body_invweight0_id, body1][1] + body_invweight0[body_invweight0_id, body2][1]
 
   for i in range(3):
@@ -1124,6 +1347,8 @@ def _equality_weld(
       efc_frictionloss_out,
     )
 
+    efc_aref_out[worldid, efcid + 3 + i] -= Jdotv_r[i]
+
 
 @wp.kernel
 def _friction_dof(
@@ -1680,6 +1905,7 @@ def _limit_tendon(
     )
 
 
+@cache_kernel
 def _efc_contact_init(cone_type: types.ConeType, is_sparse: bool):
   IS_ELLIPTIC = cone_type == types.ConeType.ELLIPTIC
   IS_SPARSE = is_sparse
@@ -1803,6 +2029,7 @@ def _efc_contact_init(cone_type: types.ConeType, is_sparse: bool):
   return kernel
 
 
+@cache_kernel
 def _efc_contact_jac_sparse(cone_type: types.ConeType):
   IS_ELLIPTIC = cone_type == types.ConeType.ELLIPTIC
 
@@ -1973,6 +2200,7 @@ def _efc_contact_jac_sparse(cone_type: types.ConeType):
   return kernel
 
 
+@cache_kernel
 def _efc_contact_jac_dense(tile_size: int, cone_type: types.ConeType):
   TILE_SIZE = tile_size
   IS_ELLIPTIC = cone_type == types.ConeType.ELLIPTIC
@@ -2117,6 +2345,7 @@ def _efc_contact_jac_dense(tile_size: int, cone_type: types.ConeType):
   return kernel
 
 
+@cache_kernel
 def _efc_contact_update(cone_type: types.ConeType):
   IS_ELLIPTIC = cone_type == types.ConeType.ELLIPTIC
 
@@ -2297,7 +2526,10 @@ def make_constraint(m: types.Model, d: types.Data):
           m.body_dofnum,
           m.body_dofadr,
           m.body_invweight0,
+          m.jnt_type,
+          m.jnt_dofadr,
           m.dof_bodyid,
+          m.dof_jntid,
           m.dof_parentid,
           m.site_bodyid,
           m.eq_obj1id,
@@ -2316,6 +2548,9 @@ def make_constraint(m: types.Model, d: types.Data):
           d.site_xpos,
           d.subtree_com,
           d.cdof,
+          d.cvel,
+          d.cdof_dot,
+          d.subtree_linvel,
           d.njmax,
           d.njmax_nnz,
         ],
@@ -2351,7 +2586,10 @@ def make_constraint(m: types.Model, d: types.Data):
           m.body_dofnum,
           m.body_dofadr,
           m.body_invweight0,
+          m.jnt_type,
+          m.jnt_dofadr,
           m.dof_bodyid,
+          m.dof_jntid,
           m.dof_parentid,
           m.site_bodyid,
           m.site_quat,
@@ -2372,6 +2610,9 @@ def make_constraint(m: types.Model, d: types.Data):
           d.site_xpos,
           d.subtree_com,
           d.cdof,
+          d.cvel,
+          d.cdof_dot,
+          d.subtree_linvel,
           d.njmax,
           d.njmax_nnz,
         ],
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/derivative.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/derivative.py
index 4a241f3b..28c8193f 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/derivative.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/derivative.py
@@ -15,6 +15,7 @@
 
 import warp as wp
 
+from mujoco.mjx.third_party.mujoco_warp._src import math
 from mujoco.mjx.third_party.mujoco_warp._src import util_misc
 from mujoco.mjx.third_party.mujoco_warp._src.support import next_act
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
@@ -24,6 +25,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import DisableBit
 from mujoco.mjx.third_party.mujoco_warp._src.types import DynType
 from mujoco.mjx.third_party.mujoco_warp._src.types import GainType
 from mujoco.mjx.third_party.mujoco_warp._src.types import Model
+from mujoco.mjx.third_party.mujoco_warp._src.types import vec10
 from mujoco.mjx.third_party.mujoco_warp._src.types import vec10f
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
 
@@ -60,13 +62,47 @@ def _qderiv_actuator_passive_vel(
   actuator_gainprm_id = worldid % actuator_gainprm.shape[0]
   actuator_biasprm_id = worldid % actuator_biasprm.shape[0]
 
+  bias = float(0.0)
+
   if actuator_gaintype[actid] == GainType.AFFINE:
     gain = actuator_gainprm[actuator_gainprm_id, actid][2]
+  elif actuator_gaintype[actid] == GainType.DCMOTOR:
+    gain = 0.0
+    dynprm = actuator_dynprm[worldid % actuator_dynprm.shape[0], actid]
+    gainprm = actuator_gainprm[actuator_gainprm_id, actid]
+    te = dynprm[0]
+
+    # controller velocity derivative: dV/dω
+    input_mode = int(gainprm[8])
+    dVdw = 0.0
+    if input_mode == 1:
+      dVdw = -gainprm[6]  # position: -kd
+    elif input_mode == 2:
+      dVdw = -gainprm[4]  # velocity: -kp
+
+    if te > 0.0:
+      # stateful current with actearly: d(K*next_act)/dω
+      # includes both back-EMF (-K) and controller (dVdw) through act_dot
+      R = wp.max(MJ_MINVAL, gainprm[0])
+      K = gainprm[1]
+      s = 1.0 - wp.exp(-opt_timestep[worldid % opt_timestep.shape[0]] / te)
+      bias += K * (dVdw - K) * s / R
+    elif dVdw != 0.0:
+      # stateless: controller terms only (back-EMF handled in bias block)
+      R = wp.max(MJ_MINVAL, gainprm[0])
+      K = gainprm[1]
+      bias += K * dVdw / R
+
+    # LuGre: force includes -sigma1*z_dot, z_dot = a*z + v
+    # d(sigma1*z_dot)/dv = sigma1*(da/dv*z + 1), ignoring higher-order da/dv*z
+    sigma1 = dynprm[6]
+    if sigma1 > 0.0:
+      bias -= sigma1
   else:
     gain = 0.0
 
   if actuator_biastype[actid] == BiasType.AFFINE:
-    bias = actuator_biasprm[actuator_biasprm_id, actid][2]
+    bias += actuator_biasprm[actuator_biasprm_id, actid][2]
   elif actuator_biastype[actid] == BiasType.DCMOTOR:
     dynprm = actuator_dynprm[worldid % actuator_dynprm.shape[0], actid]
     te = dynprm[0]
@@ -86,11 +122,7 @@ def _qderiv_actuator_passive_vel(
         Ta = dynprm[4]
         R *= 1.0 + alpha * (T + Ta - T0)
 
-      bias = -K * K / wp.max(MJ_MINVAL, R)
-    else:
-      bias = 0.0
-  else:
-    bias = 0.0
+      bias += -K * K / wp.max(MJ_MINVAL, R)
 
   if bias == 0.0 and gain == 0.0:
     vel_out[worldid, actid] = 0.0
@@ -151,15 +183,15 @@ def _qderiv_actuator_passive_actuation_dense(
   actuator_moment_in: wp.array2d[float],
   # In:
   vel_in: wp.array2d[float],
-  qMi: wp.array[int],
-  qMj: wp.array[int],
+  Mi: wp.array[int],
+  Mj: wp.array[int],
   # Out:
   qDeriv_out: wp.array3d[float],
 ):
   worldid, elemid = wp.tid()
 
-  dofiid = qMi[elemid]
-  dofjid = qMj[elemid]
+  dofiid = Mi[elemid]
+  dofjid = Mj[elemid]
   qderiv_contrib = float(0.0)
   for actid in range(nu):
     vel = vel_in[worldid, actid]
@@ -195,8 +227,7 @@ def _qderiv_actuator_passive_actuation_dense(
 @wp.kernel
 def _qderiv_actuator_passive_actuation_sparse(
   # Model:
-  M_rownnz: wp.array[int],
-  M_rowadr: wp.array[int],
+  M_elemid: wp.array2d[int],
   # Data in:
   moment_rownnz_in: wp.array2d[int],
   moment_rowadr_in: wp.array2d[int],
@@ -204,7 +235,6 @@ def _qderiv_actuator_passive_actuation_sparse(
   actuator_moment_in: wp.array2d[float],
   # In:
   vel_in: wp.array2d[float],
-  qMj: wp.array[int],
   # Out:
   qDeriv_out: wp.array3d[float],
 ):
@@ -231,19 +261,10 @@ def _qderiv_actuator_passive_actuation_sparse(
         continue
       dofj = moment_colind_in[worldid, rowadrj]
 
-      contrib = moment_i * moment_j * vel
-
-      # Search the corresponding elemid
-      # TODO: This could be precalculated for improved performance
-      row = dofi
-      col = dofj
-      row_startk = M_rowadr[row] - 1
-      row_nnz = M_rownnz[row]
-      for k in range(row_nnz):
-        row_startk += 1
-        if qMj[row_startk] == col:
-          wp.atomic_add(qDeriv_out[worldid, 0], row_startk, contrib)
-          break
+      elemid = M_elemid[dofi, dofj]
+      if elemid >= 0:
+        contrib = moment_i * moment_j * vel
+        wp.atomic_add(qDeriv_out[worldid, 0], elemid, contrib)
 
 
 @wp.kernel
@@ -254,23 +275,29 @@ def _qderiv_actuator_passive(
   dof_damping: wp.array2d[float],
   dof_dampingpoly: wp.array2d[wp.vec2],
   is_sparse: bool,
+  M_elemid: wp.array2d[int],
   # Data in:
   qvel_in: wp.array2d[float],
-  qM_in: wp.array3d[float],
+  M_in: wp.array3d[float],
   # In:
-  qMi: wp.array[int],
-  qMj: wp.array[int],
+  Mi: wp.array[int],
+  Mj: wp.array[int],
   qDeriv_in: wp.array3d[float],
   # Out:
   qDeriv_out: wp.array3d[float],
 ):
   worldid, elemid = wp.tid()
 
-  dofiid = qMi[elemid]
-  dofjid = qMj[elemid]
+  dofiid = Mi[elemid]
+  dofjid = Mj[elemid]
+
+  madr = M_elemid[dofiid, dofjid]
 
   if is_sparse:
-    qderiv = qDeriv_in[worldid, 0, elemid]
+    if madr >= 0:
+      qderiv = qDeriv_in[worldid, 0, madr]
+    else:
+      qderiv = 0.0
   else:
     qderiv = qDeriv_in[worldid, dofiid, dofjid]
 
@@ -283,12 +310,13 @@ def _qderiv_actuator_passive(
   qderiv *= opt_timestep[worldid % opt_timestep.shape[0]]
 
   if is_sparse:
-    qDeriv_out[worldid, 0, elemid] = qM_in[worldid, 0, elemid] - qderiv
+    if madr >= 0:
+      qDeriv_out[worldid, 0, madr] = M_in[worldid, 0, madr] - qderiv
   else:
-    qM = qM_in[worldid, dofiid, dofjid] - qderiv
-    qDeriv_out[worldid, dofiid, dofjid] = qM
+    M = M_in[worldid, dofiid, dofjid] - qderiv
+    qDeriv_out[worldid, dofiid, dofjid] = M
     if dofiid != dofjid:
-      qDeriv_out[worldid, dofjid, dofiid] = qM
+      qDeriv_out[worldid, dofjid, dofiid] = M
 
 
 # TODO(team): improve performance with tile operations?
@@ -303,18 +331,19 @@ def _qderiv_tendon_damping(
   tendon_damping: wp.array2d[float],
   tendon_dampingpoly: wp.array2d[wp.vec2],
   is_sparse: bool,
+  M_elemid: wp.array2d[int],
   # Data in:
   ten_J_in: wp.array2d[float],
   ten_velocity_in: wp.array2d[float],
   # In:
-  qMi: wp.array[int],
-  qMj: wp.array[int],
+  Mi: wp.array[int],
+  Mj: wp.array[int],
   # Out:
   qDeriv_out: wp.array3d[float],
 ):
   worldid, elemid = wp.tid()
-  dofiid = qMi[elemid]
-  dofjid = qMj[elemid]
+  dofiid = Mi[elemid]
+  dofjid = Mj[elemid]
 
   qderiv = float(0.0)
   tendon_damping_id = worldid % tendon_damping.shape[0]
@@ -343,14 +372,283 @@ def _qderiv_tendon_damping(
 
   qderiv *= opt_timestep[worldid % opt_timestep.shape[0]]
 
+  madr = M_elemid[dofiid, dofjid]
+
   if is_sparse:
-    qDeriv_out[worldid, 0, elemid] -= qderiv
+    if madr >= 0:
+      qDeriv_out[worldid, 0, madr] -= qderiv
   else:
     qDeriv_out[worldid, dofiid, dofjid] -= qderiv
     if dofiid != dofjid:
       qDeriv_out[worldid, dofjid, dofiid] -= qderiv
 
 
+@wp.kernel
+def deriv_rne_cvel_cdof_dot(
+  # Model:
+  body_parentid: wp.array[int],
+  body_jntnum: wp.array[int],
+  body_jntadr: wp.array[int],
+  body_dofadr: wp.array[int],
+  jnt_type: wp.array[int],
+  # Data in:
+  cdof_in: wp.array2d[wp.spatial_vector],
+  # In:
+  body_tree_: wp.array[int],
+  # Out:
+  Dcvel_out: wp.array3d[wp.spatial_vector],
+  Dcdof_dot_out: wp.array3d[wp.spatial_vector],
+):
+  """Forward pass: compute d(cvel)/d(qvel_k) and d(cdof_dot)/d(qvel_k).
+
+  Mirrors the accumulation order of comvel for each joint type.
+
+  Dcdof_dot for rotation DOFs of free joints (dofid+0..2) is zero because the
+  forward pass sets cdof_dot[dofid+0..2] = 0.  The Dcdof_dot array is
+  zero-initialized so no explicit write is needed.
+  """
+  worldid, nodeid, dofid = wp.tid()
+  bodyid = body_tree_[nodeid]
+  dofadr = body_dofadr[bodyid]
+  jntid = body_jntadr[bodyid]
+  jntnum = body_jntnum[bodyid]
+  pid = body_parentid[bodyid]
+
+  cdof = cdof_in[worldid]
+
+  # Initialize from parent
+  cvel_k = Dcvel_out[worldid, pid, dofid]
+
+  if jntnum == 0:
+    Dcvel_out[worldid, bodyid, dofid] = cvel_k
+    return
+
+  dof_i = dofadr
+
+  for j in range(jntid, jntid + jntnum):
+    jnttype = jnt_type[j]
+
+    if jnttype == 0:  # FREE
+      # rotation DOFs (dof_i+0..2) contribute to cvel
+      if dofid >= dof_i and dofid < dof_i + 3:
+        cvel_k += cdof[dofid]
+
+      # cdof_dot for rotation DOFs is zero (set in forward kinematics),
+      # so Dcdof_dot for rotation DOFs is zero (from wp.zeros init)
+
+      # derivative of cdof_dot for translation DOFs 3,4,5
+      Dcdof_dot_out[worldid, dof_i + 3, dofid] = math.motion_cross(cvel_k, cdof[dof_i + 3])
+      Dcdof_dot_out[worldid, dof_i + 4, dofid] = math.motion_cross(cvel_k, cdof[dof_i + 4])
+      Dcdof_dot_out[worldid, dof_i + 5, dofid] = math.motion_cross(cvel_k, cdof[dof_i + 5])
+
+      # translation DOFs (dof_i+3..5) contribute to cvel
+      if dofid >= dof_i + 3 and dofid < dof_i + 6:
+        cvel_k += cdof[dofid]
+
+      dof_i += 6
+
+    elif jnttype == 1:  # BALL
+      Dcdof_dot_out[worldid, dof_i + 0, dofid] = math.motion_cross(cvel_k, cdof[dof_i + 0])
+      Dcdof_dot_out[worldid, dof_i + 1, dofid] = math.motion_cross(cvel_k, cdof[dof_i + 1])
+      Dcdof_dot_out[worldid, dof_i + 2, dofid] = math.motion_cross(cvel_k, cdof[dof_i + 2])
+
+      if dofid >= dof_i and dofid < dof_i + 3:
+        cvel_k += cdof[dofid]
+
+      dof_i += 3
+    else:  # HINGE or SLIDE
+      Dcdof_dot_out[worldid, dof_i, dofid] = math.motion_cross(cvel_k, cdof[dof_i])
+
+      if dofid == dof_i:
+        cvel_k += cdof[dof_i]
+
+      dof_i += 1
+
+  Dcvel_out[worldid, bodyid, dofid] = cvel_k
+
+
+@wp.kernel
+def deriv_rne_cacc_cfrcbody_forward(
+  # Model:
+  body_parentid: wp.array[int],
+  body_dofnum: wp.array[int],
+  body_dofadr: wp.array[int],
+  # Data in:
+  qvel_in: wp.array2d[float],
+  cinert_in: wp.array2d[vec10],
+  cvel_in: wp.array2d[wp.spatial_vector],
+  cdof_dot_in: wp.array2d[wp.spatial_vector],
+  # In:
+  body_tree_: wp.array[int],
+  Dcvel_in: wp.array3d[wp.spatial_vector],
+  Dcdof_dot_in: wp.array3d[wp.spatial_vector],
+  # Out:
+  Dcacc_out: wp.array3d[wp.spatial_vector],
+  Dcfrcbody_out: wp.array3d[wp.spatial_vector],
+):
+  """Forward pass: compute d(cacc)/d(qvel_k) and d(cfrc_body)/d(qvel_k)."""
+  worldid, nodeid, dofid = wp.tid()
+  bodyid = body_tree_[nodeid]
+  dofadr = body_dofadr[bodyid]
+  dofnum = body_dofnum[bodyid]
+  pid = body_parentid[bodyid]
+
+  qvel = qvel_in[worldid]
+
+  dcacc = Dcacc_out[worldid, pid, dofid]
+
+  for j in range(dofadr, dofadr + dofnum):
+    # Term 1: d(cdof_dot * qvel)/d(qvel_k) when j == dofid
+    if j == dofid:
+      dcacc += cdof_dot_in[worldid, j]
+
+    # Term 2: cdof_dot depends on cvel which depends on qvel_k
+    dcdofdot = Dcdof_dot_in[worldid, j, dofid]
+    dcacc += dcdofdot * qvel[j]
+
+  Dcacc_out[worldid, bodyid, dofid] = dcacc
+
+  # d(cfrc_body)/d(qvel_k)
+  cinert = cinert_in[worldid, bodyid]
+  cvel = cvel_in[worldid, bodyid]
+  dcvel = Dcvel_in[worldid, bodyid, dofid]
+
+  # term1 = cinert * d(cacc)/d(qvel_k)
+  term1 = math.inert_vec(cinert, dcacc)
+
+  # term2 = d(cvel x* (cinert * cvel))/d(qvel_k)
+  cinert_cvel = math.inert_vec(cinert, cvel)
+  cinert_dcvel = math.inert_vec(cinert, dcvel)
+  term2 = math.motion_cross_force(dcvel, cinert_cvel) + math.motion_cross_force(cvel, cinert_dcvel)
+
+  Dcfrcbody_out[worldid, bodyid, dofid] = term1 + term2
+
+
+@wp.kernel
+def deriv_rne_cfrcbody_backward(
+  # Model:
+  body_parentid: wp.array[int],
+  # In:
+  body_tree_: wp.array[int],
+  # Out:
+  Dcfrcbody_out: wp.array3d[wp.spatial_vector],
+):
+  """Backward pass: accumulate d(cfrc_body) from children to parents."""
+  worldid, nodeid, dofid = wp.tid()
+  bodyid = body_tree_[nodeid]
+  pid = body_parentid[bodyid]
+
+  # body_tree never contains bodyid=0 (worldbody), so pid >= 0 is always valid.
+  # Siblings at the same level may share a parent; atomic_add handles this.
+  val = Dcfrcbody_out[worldid, bodyid, dofid]
+  wp.atomic_add(Dcfrcbody_out[worldid, pid], dofid, val)
+
+
+@wp.kernel
+def deriv_rne_body2jnt_sparse(
+  # Model:
+  dof_bodyid: wp.array[int],
+  # Data in:
+  cdof_in: wp.array2d[wp.spatial_vector],
+  # In:
+  timestep: wp.array[float],
+  Di: wp.array[int],
+  Dj: wp.array[int],
+  Dcfrcbody_in: wp.array3d[wp.spatial_vector],
+  flg_subtract: bool,
+  # Out:
+  qDeriv_out: wp.array3d[float],
+):
+  """Project body-space RNE derivatives into joint-space qDeriv (sparse)."""
+  worldid, elemid = wp.tid()
+  dt = timestep[worldid % timestep.shape[0]]
+
+  i = Di[elemid]
+  j = Dj[elemid]
+
+  body_i = dof_bodyid[i]
+  dcfrc = Dcfrcbody_in[worldid, body_i, j]
+  term = wp.dot(cdof_in[worldid, i], dcfrc)
+
+  if flg_subtract:
+    wp.atomic_sub(qDeriv_out[worldid, 0], elemid, dt * term)
+  else:
+    wp.atomic_add(qDeriv_out[worldid, 0], elemid, dt * term)
+
+
+def deriv_rne_vel(m: Model, d: Data, out: wp.array3d[float], flg_subtract: bool = False):
+  """Compute RNE velocity derivatives and add/subtract from the output.
+
+  Implements the analytical derivative of inverse-dynamics Coriolis/centrifugal
+  forces with respect to joint velocities.
+
+  Args:
+    m: The model (device).
+    d: The data (device).
+    out: D-structure output array (nworld, 1, nD) to accumulate RNE terms into.
+    flg_subtract: If True, subtract the RNE derivatives from output instead of adding them.
+  """
+  # TODO(team): consider caching these allocations
+  Dcvel = wp.zeros((d.nworld, m.nbody, m.nv), dtype=wp.spatial_vector)
+  Dcdof_dot = wp.zeros((d.nworld, m.nv, m.nv), dtype=wp.spatial_vector)
+  Dcacc = wp.zeros((d.nworld, m.nbody, m.nv), dtype=wp.spatial_vector)
+  Dcfrcbody = wp.zeros((d.nworld, m.nbody, m.nv), dtype=wp.spatial_vector)
+
+  # Forward pass 1: compute Dcvel and Dcdof_dot
+  for body_tree in m.body_tree:
+    wp.launch(
+      deriv_rne_cvel_cdof_dot,
+      dim=(d.nworld, body_tree.size, m.nv),
+      inputs=[
+        m.body_parentid,
+        m.body_jntnum,
+        m.body_jntadr,
+        m.body_dofadr,
+        m.jnt_type,
+        d.cdof,
+        body_tree,
+      ],
+      outputs=[Dcvel, Dcdof_dot],
+    )
+
+  # Forward pass 2: compute Dcacc and Dcfrcbody
+  for body_tree in m.body_tree:
+    wp.launch(
+      deriv_rne_cacc_cfrcbody_forward,
+      dim=(d.nworld, body_tree.size, m.nv),
+      inputs=[
+        m.body_parentid,
+        m.body_dofnum,
+        m.body_dofadr,
+        d.qvel,
+        d.cinert,
+        d.cvel,
+        d.cdof_dot,
+        body_tree,
+        Dcvel,
+        Dcdof_dot,
+      ],
+      outputs=[Dcacc, Dcfrcbody],
+    )
+
+  # Backward pass: accumulate Dcfrcbody from children to parents
+  for body_tree in reversed(m.body_tree):
+    wp.launch(
+      deriv_rne_cfrcbody_backward,
+      dim=(d.nworld, body_tree.size, m.nv),
+      inputs=[m.body_parentid, body_tree],
+      outputs=[Dcfrcbody],
+    )
+
+  # Project body-space derivatives into joint-space qDeriv (always sparse D-structure)
+  wp.launch(
+    deriv_rne_body2jnt_sparse,
+    dim=(d.nworld, m.qD_fullm_i.size),
+    inputs=[m.dof_bodyid, d.cdof, m.opt.timestep, m.qD_fullm_i, m.qD_fullm_j, Dcfrcbody, flg_subtract],
+    outputs=[out],
+  )
+
+
 @event_scope
 def deriv_smooth_vel(m: Model, d: Data, out: wp.array2d[float]):
   """Analytical derivative of smooth forces w.r.t. velocities.
@@ -358,12 +656,10 @@ def deriv_smooth_vel(m: Model, d: Data, out: wp.array2d[float]):
   Args:
     m: The model containing kinematic and dynamic information (device).
     d: The data object containing the current state and output arrays (device).
-    out: qM - dt * qDeriv (derivatives of smooth forces w.r.t velocities).
+    out: M - dt * qDeriv (derivatives of smooth forces w.r.t velocities).
   """
-  qMi = m.qM_fullm_i
-  qMj = m.qM_fullm_j
-
-  # TODO(team): implicit requires different sparsity structure
+  Mi = m.M_fullm_i
+  Mj = m.M_fullm_j
 
   if ~(m.opt.disableflags & (DisableBit.ACTUATION | DisableBit.DAMPER)):
     # TODO(team): only clear elements not set by _qderiv_actuator_passive
@@ -399,41 +695,49 @@ def deriv_smooth_vel(m: Model, d: Data, out: wp.array2d[float]):
         wp.launch(
           _qderiv_actuator_passive_actuation_sparse,
           dim=(d.nworld, m.nu),
-          inputs=[m.M_rownnz, m.M_rowadr, d.moment_rownnz, d.moment_rowadr, d.moment_colind, d.actuator_moment, vel, qMj],
+          inputs=[
+            m.M_elemid,
+            d.moment_rownnz,
+            d.moment_rowadr,
+            d.moment_colind,
+            d.actuator_moment,
+            vel,
+          ],
           outputs=[out],
         )
       else:
         wp.launch(
           _qderiv_actuator_passive_actuation_dense,
-          dim=(d.nworld, qMi.size),
-          inputs=[m.nu, d.moment_rownnz, d.moment_rowadr, d.moment_colind, d.actuator_moment, vel, qMi, qMj],
+          dim=(d.nworld, Mi.size),
+          inputs=[m.nu, d.moment_rownnz, d.moment_rowadr, d.moment_colind, d.actuator_moment, vel, Mi, Mj],
           outputs=[out],
         )
     wp.launch(
       _qderiv_actuator_passive,
-      dim=(d.nworld, qMi.size),
+      dim=(d.nworld, Mi.size),
       inputs=[
         m.opt.timestep,
         m.opt.disableflags,
         m.dof_damping,
         m.dof_dampingpoly,
         m.is_sparse,
+        m.M_elemid,
         d.qvel,
-        d.qM,
-        qMi,
-        qMj,
+        d.M,
+        Mi,
+        Mj,
         out,
       ],
       outputs=[out],
     )
   else:
-    # TODO(team): directly utilize qM for these settings
-    wp.copy(out, d.qM)
+    # TODO(team): directly utilize M for these settings
+    wp.copy(out, d.M)
 
   if not (m.opt.disableflags & DisableBit.DAMPER):
     wp.launch(
       _qderiv_tendon_damping,
-      dim=(d.nworld, qMi.size),
+      dim=(d.nworld, Mi.size),
       inputs=[
         m.ntendon,
         m.opt.timestep,
@@ -443,12 +747,11 @@ def deriv_smooth_vel(m: Model, d: Data, out: wp.array2d[float]):
         m.tendon_damping,
         m.tendon_dampingpoly,
         m.is_sparse,
+        m.M_elemid,
         d.ten_J,
         d.ten_velocity,
-        qMi,
-        qMj,
+        Mi,
+        Mj,
       ],
       outputs=[out],
     )
-
-  # TODO(team): rne derivative
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py
index c9b3c153..00f22815 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py
@@ -20,12 +20,14 @@ import warp as wp
 from mujoco.mjx.third_party.mujoco_warp._src import collision_driver
 from mujoco.mjx.third_party.mujoco_warp._src import constraint
 from mujoco.mjx.third_party.mujoco_warp._src import derivative
+from mujoco.mjx.third_party.mujoco_warp._src import history
 from mujoco.mjx.third_party.mujoco_warp._src import island
 from mujoco.mjx.third_party.mujoco_warp._src import math
 from mujoco.mjx.third_party.mujoco_warp._src import passive
 from mujoco.mjx.third_party.mujoco_warp._src import sensor
 from mujoco.mjx.third_party.mujoco_warp._src import smooth
 from mujoco.mjx.third_party.mujoco_warp._src import solver
+from mujoco.mjx.third_party.mujoco_warp._src import types
 from mujoco.mjx.third_party.mujoco_warp._src import util_misc
 from mujoco.mjx.third_party.mujoco_warp._src.support import next_act
 from mujoco.mjx.third_party.mujoco_warp._src.support import xfrc_accumulate
@@ -302,6 +304,9 @@ def _advance(m: Model, d: Data, qacc: wp.array, qvel: Optional[wp.array] = None)
     outputs=[d.qpos],
   )
 
+  # advance history buffers before time advance
+  history.insert_ctrl_history(m, d)
+
   wp.launch(
     _next_time,
     dim=d.nworld,
@@ -346,17 +351,18 @@ def _compute_damping_deriv(
 def _euler_damp_qfrc_sparse(
   # Model:
   opt_timestep: wp.array[float],
-  dof_Madr: wp.array[int],
+  M_rownnz: wp.array[int],
+  M_rowadr: wp.array[int],
   # In:
   damp_deriv: wp.array2d[float],
   # Out:
-  qM_integration_out: wp.array3d[float],
+  M_integration_out: wp.array3d[float],
 ):
   worldid, tid = wp.tid()
   timestep = opt_timestep[worldid % opt_timestep.shape[0]]
 
-  adr = dof_Madr[tid]
-  qM_integration_out[worldid, 0, adr] += timestep * damp_deriv[worldid, tid]
+  adr = M_rowadr[tid] + M_rownnz[tid] - 1
+  M_integration_out[worldid, 0, adr] += timestep * damp_deriv[worldid, tid]
 
 
 @cache_kernel
@@ -366,7 +372,7 @@ def _tile_euler_dense(tile: TileSet):
     # Model:
     opt_timestep: wp.array[float],
     # Data in:
-    qM_in: wp.array3d[float],
+    M_in: wp.array3d[float],
     efc_Ma_in: wp.array2d[float],
     # In:
     damp_deriv: wp.array2d[float],
@@ -379,14 +385,14 @@ def _tile_euler_dense(tile: TileSet):
     TILE_SIZE = wp.static(tile.size)
 
     dofid = adr_in[nodeid]
-    M_tile = wp.tile_load(qM_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid))
+    M_tile = wp.tile_load(M_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid))
     damping_tile = wp.tile_load(damp_deriv[worldid], shape=(TILE_SIZE,), offset=(dofid,))
     damping_scaled = damping_tile * timestep
     qm_integration_tile = wp.tile_diag_add(M_tile, damping_scaled)
 
     Ma_tile = wp.tile_load(efc_Ma_in[worldid], shape=(TILE_SIZE,), offset=(dofid,))
-    L_tile = wp.tile_cholesky(qm_integration_tile)
-    qacc_tile = wp.tile_cholesky_solve(L_tile, Ma_tile)
+    L_tile = wp.tile_cholesky(qm_integration_tile, fill_mode="upper")
+    qacc_tile = wp.tile_cholesky_solve(L_tile, Ma_tile, fill_mode="upper")
     wp.tile_store(qacc_out[worldid], qacc_tile, offset=(dofid))
 
   return euler_dense
@@ -409,22 +415,22 @@ def euler(m: Model, d: Data):
     )
 
     if m.is_sparse:
-      qM = wp.clone(d.qM)
+      M = wp.clone(d.M)
       qLD = wp.empty((d.nworld, 1, m.nC), dtype=float)
       qLDiagInv = wp.empty((d.nworld, m.nv), dtype=float)
       wp.launch(
         _euler_damp_qfrc_sparse,
         dim=(d.nworld, m.nv),
-        inputs=[m.opt.timestep, m.dof_Madr, damp_deriv],
-        outputs=[qM],
+        inputs=[m.opt.timestep, m.M_rownnz, m.M_rowadr, damp_deriv],
+        outputs=[M],
       )
-      smooth.factor_solve_i(m, d, qM, qLD, qLDiagInv, qacc, d.efc.Ma)
+      smooth.factor_solve_i(m, d, M, qLD, qLDiagInv, qacc, d.efc.Ma)
     else:
-      for tile in m.qM_tiles:
+      for tile in m.M_tiles:
         wp.launch_tiled(
           _tile_euler_dense(tile),
           dim=(d.nworld, tile.adr.size),
-          inputs=[m.opt.timestep, d.qM, d.efc.Ma, damp_deriv, tile.adr],
+          inputs=[m.opt.timestep, d.M, d.efc.Ma, damp_deriv, tile.adr],
           outputs=[qacc],
           block_dim=m.block_dim.euler_dense,
         )
@@ -573,16 +579,62 @@ def rungekutta4(m: Model, d: Data):
   _advance(m, d, qacc_rk, qvel_rk)
 
 
+@wp.kernel
+def _map_m2d(
+  # Model:
+  mapM2D: wp.array[int],
+  is_sparse: bool,
+  # In:
+  qDi: wp.array[int],
+  qDj: wp.array[int],
+  qH_M: wp.array3d[float],
+  # Data out:
+  qLU_out: wp.array3d[float],
+):
+  worldid, elemid = wp.tid()
+  if is_sparse:
+    m_idx = mapM2D[elemid]
+    if m_idx >= 0:
+      qLU_out[worldid, 0, elemid] = qH_M[worldid, 0, m_idx]
+    else:
+      qLU_out[worldid, 0, elemid] = 0.0
+  else:
+    i = qDi[elemid]
+    j = qDj[elemid]
+    qLU_out[worldid, 0, elemid] = qH_M[worldid, i, j]
+
+
 @event_scope
 def implicit(m: Model, d: Data):
   """Integrates fully implicit in velocity."""
-  if ~(m.opt.disableflags | ~(DisableBit.ACTUATION | DisableBit.SPRING | DisableBit.DAMPER)):
+  if m.opt.integrator == IntegratorType.IMPLICIT:
+    qH_M = wp.empty(d.M.shape, dtype=float)
+
+    # 1. Compute M - dt * qDeriv_smooth in M-structure
+    derivative.deriv_smooth_vel(m, d, qH_M)
+
+    # 2. Map M-structure to D-structure
+    wp.launch(
+      _map_m2d,
+      dim=(d.nworld, m.nD),
+      inputs=[m.mapM2D, m.is_sparse, m.qD_fullm_i, m.qD_fullm_j, qH_M],
+      outputs=[d.qLU],
+    )
+
+    # 3. Compute RNE derivatives, scale by timestep, and subtract in-place from qLU
+    derivative.deriv_rne_vel(m, d, d.qLU, flg_subtract=True)
+
+    # 4. Factorize and solve: qacc = qLU \ Ma
+    qacc = wp.empty((d.nworld, m.nv), dtype=float)
+    smooth.factor_solve_lu(m, d, d.qLU, qacc, d.efc.Ma)
+    _advance(m, d, qacc)
+  elif ~(m.opt.disableflags | ~(DisableBit.ACTUATION | DisableBit.SPRING | DisableBit.DAMPER)):
     if m.is_sparse:
-      qDeriv = wp.empty((d.nworld, 1, m.nM), dtype=float)
+      qDeriv = wp.empty((d.nworld, 1, m.nC), dtype=float)
       qLD = wp.empty((d.nworld, 1, m.nC), dtype=float)
     else:
-      qDeriv = wp.empty(d.qM.shape, dtype=float)
-      qLD = wp.empty(d.qM.shape, dtype=float)
+      qDeriv = wp.empty(d.M.shape, dtype=float)
+      qLD = wp.empty(d.M.shape, dtype=float)
     qLDiagInv = wp.empty((d.nworld, m.nv), dtype=float)
     derivative.deriv_smooth_vel(m, d, qDeriv)
     qacc = wp.empty((d.nworld, m.nv), dtype=float)
@@ -613,8 +665,7 @@ def fwd_position(m: Model, d: Data, factorize: bool = True):
   if m.opt.run_collision_detection:
     collision_driver.collision(m, d)
   constraint.make_constraint(m, d)
-  # TODO(team): remove False after island features are more complete
-  if False and not (m.opt.disableflags & DisableBit.ISLAND):
+  if m.ntree > 1 and not (m.opt.disableflags & types.DisableBit.ISLAND):
     island.island(m, d)
   smooth.transmission(m, d)
 
@@ -1098,6 +1149,13 @@ def fwd_actuation(m: Model, d: Data):
     d.actuator_force.zero_()
     return
 
+  # read delayed ctrl (or direct copy if no delay)
+  if m.nhistory > 0:
+    ctrl = wp.empty((d.nworld, m.nu), dtype=float)
+    history.read_ctrl_delayed(m, d, ctrl)
+  else:
+    ctrl = d.ctrl
+
   wp.launch(
     _actuator_force,
     dim=(d.nworld, m.nu),
@@ -1122,7 +1180,7 @@ def fwd_actuation(m: Model, d: Data):
       m.actuator_acc0,
       m.actuator_lengthrange,
       d.act,
-      d.ctrl,
+      ctrl,
       d.actuator_length,
       d.actuator_velocity,
       m.opt.disableflags & DisableBit.CLAMPCTRL,
@@ -1221,7 +1279,7 @@ def fwd_acceleration(m: Model, d: Data, factorize: bool = False):
   xfrc_accumulate(m, d, d.qfrc_smooth)
 
   if factorize:
-    smooth.factor_solve_i(m, d, d.qM, d.qLD, d.qLDiagInv, d.qacc_smooth, d.qfrc_smooth)
+    smooth.factor_solve_i(m, d, d.M, d.qLD, d.qLDiagInv, d.qacc_smooth, d.qfrc_smooth)
   else:
     smooth.solve_m(m, d, d.qacc_smooth, d.qfrc_smooth)
 
@@ -1266,7 +1324,7 @@ def step(m: Model, d: Data):
     euler(m, d)
   elif m.opt.integrator == IntegratorType.RK4:
     rungekutta4(m, d)
-  elif m.opt.integrator == IntegratorType.IMPLICITFAST:
+  elif m.opt.integrator in (IntegratorType.IMPLICITFAST, IntegratorType.IMPLICIT):
     implicit(m, d)
   else:
     raise NotImplementedError(f"integrator {m.opt.integrator} not implemented.")
@@ -1307,8 +1365,7 @@ def step2(m: Model, d: Data):
   sensor.sensor_acc(m, d)
 
   # integrate with Euler or implicitfast
-  # TODO(team): implicit
-  if m.opt.integrator == IntegratorType.IMPLICITFAST:
+  if m.opt.integrator in (IntegratorType.IMPLICITFAST, IntegratorType.IMPLICIT):
     implicit(m, d)
   else:
     # note: RK4 defaults to Euler
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/history.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/history.py
new file mode 100644
index 00000000..42149578
--- /dev/null
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/history.py
@@ -0,0 +1,925 @@
+# Copyright 2026 The Newton Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import warp as wp
+
+from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
+from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
+from mujoco.mjx.third_party.mujoco_warp._src.types import Data
+from mujoco.mjx.third_party.mujoco_warp._src.types import Model
+
+wp.set_module_options({"enable_backward": False})
+
+
+@wp.func
+def _history_physical_index(cursor: int, n: int, logical: int) -> int:
+  """Convert logical index (0=oldest, n-1=newest) to physical index."""
+  return (cursor + 1 + logical) % n
+
+
+@wp.func
+def _history_find_index(
+  # In:
+  buf: wp.array2d[float],
+  worldid: int,
+  buf_offset: int,
+  n: int,
+  cursor: int,
+  t: float,
+) -> int:
+  """Find logical index i such that times[i-1] < t <= times[i].
+
+  Returns 0 if t <= times[oldest], n if t > times[newest].
+  Uses circular binary search matching MuJoCo C historyFindIndex.
+  """
+  times_offset = buf_offset + 2
+
+  oldest_phys = _history_physical_index(cursor, n, 0)
+  newest_phys = _history_physical_index(cursor, n, n - 1)
+  t_oldest = buf[worldid, times_offset + oldest_phys]
+  t_newest = buf[worldid, times_offset + newest_phys]
+
+  # before or at first element
+  if t <= t_oldest:
+    return 0
+
+  # after last element
+  if t > t_newest:
+    return n
+
+  # circular binary search: find smallest logical i such that times[phys(i)] >= t
+  lo = int(0)
+  hi = int(n - 1)
+  while hi - lo > 1:
+    mid = int((lo + hi) >> 1)
+    mid_phys = _history_physical_index(cursor, n, mid)
+    if buf[worldid, times_offset + mid_phys] < t:
+      lo = mid
+    else:
+      hi = mid
+
+  return hi
+
+
+@wp.func
+def _history_read_scalar(
+  # In:
+  buf: wp.array2d[float],
+  worldid: int,
+  buf_offset: int,
+  n: int,
+  t: float,
+  interp: int,
+) -> float:
+  """Read a scalar value from history buffer at time t.
+
+  interp: 0=zero-order-hold, 1=linear, 2=cubic (Catmull-Rom spline)
+  """
+  cursor = int(buf[worldid, buf_offset + 1])
+  times_offset = buf_offset + 2
+  values_offset = buf_offset + 2 + n
+
+  oldest_phys = _history_physical_index(cursor, n, 0)
+  newest_phys = _history_physical_index(cursor, n, n - 1)
+  t_oldest = buf[worldid, times_offset + oldest_phys]
+  t_newest = buf[worldid, times_offset + newest_phys]
+
+  # extrapolate before oldest
+  if t <= t_oldest + 1e-6:
+    return buf[worldid, values_offset + oldest_phys]
+
+  # extrapolate after newest
+  if t >= t_newest - 1e-6:
+    return buf[worldid, values_offset + newest_phys]
+
+  # find bracketing index
+  i = _history_find_index(buf, worldid, buf_offset, n, cursor, t)
+  phys_i = _history_physical_index(cursor, n, i)
+
+  # exact match
+  if wp.abs(t - buf[worldid, times_offset + phys_i]) < 1e-6:
+    return buf[worldid, values_offset + phys_i]
+
+  phys_lo = _history_physical_index(cursor, n, i - 1)
+  phys_hi = phys_i
+
+  # zero-order hold
+  if interp == 0:
+    return buf[worldid, values_offset + phys_lo]
+
+  dt = buf[worldid, times_offset + phys_hi] - buf[worldid, times_offset + phys_lo]
+  alpha = (t - buf[worldid, times_offset + phys_lo]) / dt
+  v_lo = buf[worldid, values_offset + phys_lo]
+  v_hi = buf[worldid, values_offset + phys_hi]
+
+  # linear interpolation
+  if interp == 1:
+    return v_lo + alpha * (v_hi - v_lo)
+
+  # cubic spline interpolation (Catmull-Rom)
+  alpha2 = alpha * alpha
+  alpha3 = alpha2 * alpha
+  h00 = 2.0 * alpha3 - 3.0 * alpha2 + 1.0
+  h10 = alpha3 - 2.0 * alpha2 + alpha
+  h01 = -2.0 * alpha3 + 3.0 * alpha2
+  h11 = alpha3 - alpha2
+
+  # finite-differenced Catmull-Rom slopes, 0 at endpoints
+  m_lo = 0.0
+  if i > 1:
+    phys_lo_prev = _history_physical_index(cursor, n, i - 2)
+    dt_lo = buf[worldid, times_offset + phys_hi] - buf[worldid, times_offset + phys_lo_prev]
+    m_lo = (v_hi - buf[worldid, values_offset + phys_lo_prev]) / dt_lo
+
+  m_hi = 0.0
+  if i < n - 1:
+    phys_hi_next = _history_physical_index(cursor, n, i + 1)
+    dt_hi = buf[worldid, times_offset + phys_hi_next] - buf[worldid, times_offset + phys_lo]
+    m_hi = (buf[worldid, values_offset + phys_hi_next] - v_lo) / dt_hi
+
+  return h00 * v_lo + h10 * dt * m_lo + h01 * v_hi + h11 * dt * m_hi
+
+
+@wp.func
+def _history_read_vector(
+  # In:
+  adr: int,
+  buf: wp.array2d[float],
+  worldid: int,
+  buf_offset: int,
+  n: int,
+  dim: int,
+  t: float,
+  interp: int,
+  # Data out:
+  sensordata_out: wp.array2d[float],
+) -> int:
+  """Read a vector value from history buffer at time t into sensordata.
+
+  Returns 1 on success (value written to sensordata).
+  interp: 0=zero-order-hold, 1=linear, 2=cubic (Catmull-Rom spline)
+  """
+  cursor = int(buf[worldid, buf_offset + 1])
+  times_offset = buf_offset + 2
+  values_offset = buf_offset + 2 + n
+
+  oldest_phys = _history_physical_index(cursor, n, 0)
+  newest_phys = _history_physical_index(cursor, n, n - 1)
+  t_oldest = buf[worldid, times_offset + oldest_phys]
+  t_newest = buf[worldid, times_offset + newest_phys]
+
+  # extrapolate before oldest: copy oldest
+  if t <= t_oldest + 1e-6:
+    for d in range(dim):
+      sensordata_out[worldid, adr + d] = buf[worldid, values_offset + oldest_phys * dim + d]
+    return 1
+
+  # extrapolate after newest: copy newest
+  if t >= t_newest - 1e-6:
+    for d in range(dim):
+      sensordata_out[worldid, adr + d] = buf[worldid, values_offset + newest_phys * dim + d]
+    return 1
+
+  # find bracketing index
+  i = _history_find_index(buf, worldid, buf_offset, n, cursor, t)
+  phys_i = _history_physical_index(cursor, n, i)
+
+  # exact match
+  if wp.abs(t - buf[worldid, times_offset + phys_i]) < 1e-6:
+    for d in range(dim):
+      sensordata_out[worldid, adr + d] = buf[worldid, values_offset + phys_i * dim + d]
+    return 1
+
+  phys_lo = _history_physical_index(cursor, n, i - 1)
+  phys_hi = phys_i
+
+  # zero-order hold
+  if interp == 0:
+    for d in range(dim):
+      sensordata_out[worldid, adr + d] = buf[worldid, values_offset + phys_lo * dim + d]
+    return 1
+
+  dt = buf[worldid, times_offset + phys_hi] - buf[worldid, times_offset + phys_lo]
+  alpha = (t - buf[worldid, times_offset + phys_lo]) / dt
+
+  # linear interpolation
+  if interp == 1:
+    for d in range(dim):
+      v_lo = buf[worldid, values_offset + phys_lo * dim + d]
+      v_hi = buf[worldid, values_offset + phys_hi * dim + d]
+      sensordata_out[worldid, adr + d] = v_lo + alpha * (v_hi - v_lo)
+    return 1
+
+  # cubic spline interpolation (Catmull-Rom)
+  alpha2 = alpha * alpha
+  alpha3 = alpha2 * alpha
+  h00 = 2.0 * alpha3 - 3.0 * alpha2 + 1.0
+  h10 = alpha3 - 2.0 * alpha2 + alpha
+  h01 = -2.0 * alpha3 + 3.0 * alpha2
+  h11 = alpha3 - alpha2
+
+  for d in range(dim):
+    v_lo = buf[worldid, values_offset + phys_lo * dim + d]
+    v_hi = buf[worldid, values_offset + phys_hi * dim + d]
+
+    # finite-differenced Catmull-Rom slopes, 0 at endpoints
+    m_lo = 0.0
+    if i > 1:
+      phys_lo_prev = _history_physical_index(cursor, n, i - 2)
+      dt_lo = buf[worldid, times_offset + phys_hi] - buf[worldid, times_offset + phys_lo_prev]
+      m_lo = (v_hi - buf[worldid, values_offset + phys_lo_prev * dim + d]) / dt_lo
+
+    m_hi = 0.0
+    if i < n - 1:
+      phys_hi_next = _history_physical_index(cursor, n, i + 1)
+      dt_hi = buf[worldid, times_offset + phys_hi_next] - buf[worldid, times_offset + phys_lo]
+      m_hi = (buf[worldid, values_offset + phys_hi_next * dim + d] - v_lo) / dt_hi
+
+    sensordata_out[worldid, adr + d] = h00 * v_lo + h10 * dt * m_lo + h01 * v_hi + h11 * dt * m_hi
+  return 1
+
+
+@wp.func
+def _history_insert_scalar(
+  # In:
+  worldid: int,
+  buf_offset: int,
+  n: int,
+  t: float,
+  value: float,
+  # Out:
+  buf_out: wp.array2d[float],
+):
+  """Insert a scalar value into history buffer at time t."""
+  cursor = int(buf_out[worldid, buf_offset + 1])
+  times_offset = buf_offset + 2
+  values_offset = buf_offset + 2 + n
+
+  i = _history_find_index(buf_out, worldid, buf_offset, n, cursor, t)
+
+  # exact match
+  if i < n:
+    phys_i = _history_physical_index(cursor, n, i)
+    if wp.abs(t - buf_out[worldid, times_offset + phys_i]) < 1e-6:
+      buf_out[worldid, values_offset + phys_i] = value
+      return
+
+  # older than oldest: replace oldest
+  if i == 0:
+    oldest_phys = _history_physical_index(cursor, n, 0)
+    buf_out[worldid, times_offset + oldest_phys] = t
+    buf_out[worldid, values_offset + oldest_phys] = value
+    return
+
+  # newer than newest: advance cursor
+  if i == n:
+    cursor = (cursor + 1) % n
+    buf_out[worldid, buf_offset + 1] = float(cursor)
+    buf_out[worldid, times_offset + cursor] = t
+    buf_out[worldid, values_offset + cursor] = value
+    return
+
+  # out-of-order: shift [1, i-1] left, insert at i-1
+  for j in range(i - 1):
+    src_phys = _history_physical_index(cursor, n, j + 1)
+    dst_phys = _history_physical_index(cursor, n, j)
+    buf_out[worldid, times_offset + dst_phys] = buf_out[worldid, times_offset + src_phys]
+    buf_out[worldid, values_offset + dst_phys] = buf_out[worldid, values_offset + src_phys]
+  insert_phys = _history_physical_index(cursor, n, i - 1)
+  buf_out[worldid, times_offset + insert_phys] = t
+  buf_out[worldid, values_offset + insert_phys] = value
+
+
+@wp.func
+def _history_insert_vector(
+  # In:
+  worldid: int,
+  buf_offset: int,
+  n: int,
+  dim: int,
+  t: float,
+  src: wp.array2d[float],
+  src_adr: int,
+  # Out:
+  buf_out: wp.array2d[float],
+):
+  """Insert a vector value from src[worldid, src_adr:src_adr+dim] into history buffer at time t."""
+  cursor = int(buf_out[worldid, buf_offset + 1])
+  times_offset = buf_offset + 2
+  values_offset = buf_offset + 2 + n
+
+  i = _history_find_index(buf_out, worldid, buf_offset, n, cursor, t)
+
+  slot_phys = -1
+
+  # exact match
+  if i < n:
+    phys_i = _history_physical_index(cursor, n, i)
+    if wp.abs(t - buf_out[worldid, times_offset + phys_i]) < 1e-6:
+      slot_phys = phys_i
+
+  if slot_phys < 0:
+    if i == 0:
+      # older than oldest: replace oldest
+      slot_phys = _history_physical_index(cursor, n, 0)
+      buf_out[worldid, times_offset + slot_phys] = t
+    elif i == n:
+      # newer than newest: advance cursor
+      cursor = (cursor + 1) % n
+      buf_out[worldid, buf_offset + 1] = float(cursor)
+      slot_phys = cursor
+      buf_out[worldid, times_offset + slot_phys] = t
+    else:
+      # out-of-order: shift [1, i-1] left, insert at i-1
+      for j in range(i - 1):
+        src_phys = _history_physical_index(cursor, n, j + 1)
+        dst_phys = _history_physical_index(cursor, n, j)
+        buf_out[worldid, times_offset + dst_phys] = buf_out[worldid, times_offset + src_phys]
+        for d in range(dim):
+          buf_out[worldid, values_offset + dst_phys * dim + d] = buf_out[worldid, values_offset + src_phys * dim + d]
+      slot_phys = _history_physical_index(cursor, n, i - 1)
+      buf_out[worldid, times_offset + slot_phys] = t
+
+  # copy values
+  for d in range(dim):
+    buf_out[worldid, values_offset + slot_phys * dim + d] = src[worldid, src_adr + d]
+
+
+@wp.kernel
+def _read_ctrl_delayed_kernel(
+  # Model:
+  actuator_history: wp.array[wp.vec2i],
+  actuator_historyadr: wp.array[int],
+  actuator_delay: wp.array[float],
+  # Data in:
+  time_in: wp.array[float],
+  history_in: wp.array2d[float],
+  ctrl_in: wp.array2d[float],
+  # Data out:
+  ctrl_out: wp.array2d[float],
+):
+  """Read delayed ctrl for each actuator."""
+  worldid, uid = wp.tid()
+
+  hist = actuator_history[uid]
+  nsample = hist[0]
+  delay = actuator_delay[uid]
+
+  if nsample == 0 or delay == 0.0:
+    # no delay: direct copy
+    ctrl_out[worldid, uid] = ctrl_in[worldid, uid]
+  else:
+    interp = hist[1]
+    buf_offset = actuator_historyadr[uid]
+    t = time_in[worldid] - delay
+    ctrl_out[worldid, uid] = _history_read_scalar(history_in, worldid, buf_offset, nsample, t, interp)
+
+
+@wp.kernel
+def _insert_ctrl_history_kernel(
+  # Model:
+  actuator_history: wp.array[wp.vec2i],
+  actuator_historyadr: wp.array[int],
+  # Data in:
+  time_in: wp.array[float],
+  ctrl_in: wp.array2d[float],
+  # Data out:
+  history_out: wp.array2d[float],
+):
+  """Insert current ctrl into history buffers."""
+  worldid, uid = wp.tid()
+
+  hist = actuator_history[uid]
+  nsample = hist[0]
+  if nsample == 0:
+    return
+
+  buf_offset = actuator_historyadr[uid]
+  t = time_in[worldid]
+  value = ctrl_in[worldid, uid]
+  _history_insert_scalar(worldid, buf_offset, nsample, t, value, history_out)
+
+
+@wp.kernel
+def _insert_sensor_history_stage(
+  # Model:
+  sensor_dim: wp.array[int],
+  sensor_adr: wp.array[int],
+  sensor_history: wp.array[wp.vec2i],
+  sensor_historyadr: wp.array[int],
+  sensor_delay: wp.array[float],
+  sensor_interval: wp.array[wp.vec2],
+  # Data in:
+  time_in: wp.array[float],
+  sensordata_in: wp.array2d[float],
+  # In:
+  sensor_ids: wp.array[int],
+  # Data out:
+  history_out: wp.array2d[float],
+):
+  """Insert current sensor values into history buffers for specific sensor IDs."""
+  worldid, idx = wp.tid()
+  sid = sensor_ids[idx]
+
+  hist = sensor_history[sid]
+  nsample = hist[0]
+  if nsample == 0:
+    return
+
+  buf_offset = sensor_historyadr[sid]
+  dim = sensor_dim[sid]
+  interval_val = sensor_interval[sid]
+  period = interval_val[0]
+  t = time_in[worldid]
+
+  if period > 0.0:
+    # interval mode: check if condition is satisfied
+    time_prev = history_out[worldid, buf_offset]  # user slot stores time_prev
+    if time_prev + period <= t:
+      # advance time_prev by exact period
+      history_out[worldid, buf_offset] = time_prev + period
+      # insert sensor value
+      _history_insert_vector(worldid, buf_offset, nsample, dim, t, sensordata_in, sensor_adr[sid], history_out)
+  else:
+    _history_insert_vector(worldid, buf_offset, nsample, dim, t, sensordata_in, sensor_adr[sid], history_out)
+
+
+@wp.kernel
+def _apply_sensor_delay_kernel(
+  # Model:
+  sensor_dim: wp.array[int],
+  sensor_adr: wp.array[int],
+  sensor_history: wp.array[wp.vec2i],
+  sensor_historyadr: wp.array[int],
+  sensor_delay: wp.array[float],
+  sensor_interval: wp.array[wp.vec2],
+  # Data in:
+  time_in: wp.array[float],
+  history_in: wp.array2d[float],
+  # In:
+  sensor_ids: wp.array[int],
+  # Data out:
+  sensordata_out: wp.array2d[float],
+):
+  """Apply delay/interval logic for sensors after computation.
+
+  TODO(team): Revisit always-compute decision for computationally expensive sensors
+  with interval/period (e.g., raytracers)
+  """
+  worldid, idx = wp.tid()
+  sid = sensor_ids[idx]
+
+  hist = sensor_history[sid]
+  nsample = hist[0]
+  if nsample <= 0:
+    return
+
+  delay = sensor_delay[sid]
+  dim = sensor_dim[sid]
+  interp = hist[1]
+  buf_offset = sensor_historyadr[sid]
+  t = time_in[worldid]
+
+  if delay > 0.0:
+    # delay > 0: read delayed value from buffer
+    _history_read_vector(sensor_adr[sid], history_in, worldid, buf_offset, nsample, dim, t - delay, interp, sensordata_out)
+  else:
+    # interval-only (delay == 0, interval > 0): check interval condition
+    interval_val = sensor_interval[sid]
+    period = interval_val[0]
+    if period > 0.0:
+      time_prev = history_in[worldid, buf_offset]  # user slot
+      if time_prev + period > t:
+        # interval condition not satisfied: read from buffer
+        _history_read_vector(sensor_adr[sid], history_in, worldid, buf_offset, nsample, dim, t, interp, sensordata_out)
+      # else: interval condition satisfied, keep computed value
+
+
+def read_ctrl_delayed(m: Model, d: Data, ctrl: wp.array2d[float]):
+  """Read delayed ctrl values for all actuators."""
+  if m.nhistory == 0:
+    wp.copy(ctrl, d.ctrl)
+    return
+
+  wp.launch(
+    _read_ctrl_delayed_kernel,
+    dim=(d.nworld, m.nu),
+    inputs=[
+      m.actuator_history,
+      m.actuator_historyadr,
+      m.actuator_delay,
+      d.time,
+      d.history,
+      d.ctrl,
+    ],
+    outputs=[ctrl],
+  )
+
+
+def insert_ctrl_history(m: Model, d: Data):
+  """Insert current ctrl values into history buffers."""
+  if m.nhistory == 0 or m.nu == 0:
+    return
+
+  wp.launch(
+    _insert_ctrl_history_kernel,
+    dim=(d.nworld, m.nu),
+    inputs=[
+      m.actuator_history,
+      m.actuator_historyadr,
+      d.time,
+      d.ctrl,
+    ],
+    outputs=[d.history],
+  )
+
+
+def apply_sensor_delay(m: Model, d: Data, sensorid: wp.array[int]):
+  """Apply delay/interval logic for given sensors after computation.
+
+  Matches MuJoCo C architecture where the delayed read (mj_sensorPos) occurs
+  before the fresh value insert (mj_advance). We save fresh sensordata,
+  overwrite with delayed values, then insert the saved fresh values.
+  """
+  if m.nhistory == 0 or sensorid.shape[0] == 0:
+    return
+
+  # Save fresh sensordata before delay overwrite
+  fresh_sensordata = wp.empty_like(d.sensordata)
+  wp.copy(fresh_sensordata, d.sensordata)
+
+  # Read delayed values from buffer → overwrite sensordata
+  wp.launch(
+    _apply_sensor_delay_kernel,
+    dim=(d.nworld, sensorid.shape[0]),
+    inputs=[
+      m.sensor_dim,
+      m.sensor_adr,
+      m.sensor_history,
+      m.sensor_historyadr,
+      m.sensor_delay,
+      m.sensor_interval,
+      d.time,
+      d.history,
+      sensorid,
+    ],
+    outputs=[d.sensordata],
+  )
+
+  # Insert saved fresh sensor values into history buffers
+  wp.launch(
+    _insert_sensor_history_stage,
+    dim=(d.nworld, sensorid.shape[0]),
+    inputs=[
+      m.sensor_dim,
+      m.sensor_adr,
+      m.sensor_history,
+      m.sensor_historyadr,
+      m.sensor_delay,
+      m.sensor_interval,
+      d.time,
+      fresh_sensordata,
+      sensorid,
+    ],
+    outputs=[d.history],
+  )
+
+
+@wp.kernel
+def _read_ctrl_kernel(
+  # Model:
+  actuator_history: wp.array[wp.vec2i],
+  actuator_historyadr: wp.array[int],
+  actuator_delay: wp.array[float],
+  # Data in:
+  time_in: wp.array[float],
+  history_in: wp.array2d[float],
+  ctrl_in: wp.array2d[float],
+  # In:
+  uid: int,
+  interp: int,
+  # Out:
+  result_out: wp.array[float],
+):
+  """Read delayed ctrl for 1 actuator across all worlds."""
+  worldid = wp.tid()
+
+  hist = actuator_history[uid]
+  nsample = hist[0]
+
+  if nsample == 0:
+    result_out[worldid] = ctrl_in[worldid, uid]
+  else:
+    interp_val = interp
+    if interp_val < 0:
+      interp_val = hist[1]
+    delay = actuator_delay[uid]
+    buf_offset = actuator_historyadr[uid]
+    t = time_in[worldid] - delay
+    result_out[worldid] = _history_read_scalar(history_in, worldid, buf_offset, nsample, t, interp_val)
+
+
+def read_ctrl(
+  m: Model,
+  d: Data,
+  ctrlid: int,
+  time: wp.array[float],
+  interp: int,
+  result: wp.array2d[float],
+):
+  """Read delayed ctrl for 1 actuator across all worlds.
+
+  Args:
+    m: The model containing kinematic and dynamic information.
+    d: The data object containing the current state and output arrays.
+    ctrlid: actuator index.
+    time: query time per world (nworld,).
+    interp: interpolation mode (-1=model default, 0=ZOH, 1=linear, 2=cubic).
+    result: output buffer (nworld,).
+  """
+  wp.launch(
+    _read_ctrl_kernel,
+    dim=(d.nworld,),
+    inputs=[
+      m.actuator_history,
+      m.actuator_historyadr,
+      m.actuator_delay,
+      time,
+      d.history,
+      d.ctrl,
+      ctrlid,
+      interp,
+    ],
+    outputs=[result],
+  )
+
+
+@wp.kernel
+def _read_sensor_kernel(
+  # Model:
+  sensor_dim: wp.array[int],
+  sensor_adr: wp.array[int],
+  sensor_history: wp.array[wp.vec2i],
+  sensor_historyadr: wp.array[int],
+  sensor_delay: wp.array[float],
+  # Data in:
+  time_in: wp.array[float],
+  history_in: wp.array2d[float],
+  sensordata_in: wp.array2d[float],
+  # In:
+  sid: int,
+  interp: int,
+  # Out:
+  result_out: wp.array2d[float],
+):
+  """Read delayed sensor for 1 sensor across all worlds."""
+  worldid = wp.tid()
+
+  hist = sensor_history[sid]
+  nsample = hist[0]
+  dim = sensor_dim[sid]
+  adr = sensor_adr[sid]
+
+  if nsample == 0:
+    for i in range(dim):
+      result_out[worldid, i] = sensordata_in[worldid, adr + i]
+  else:
+    interp_val = interp
+    if interp_val < 0:
+      interp_val = hist[1]
+    delay = sensor_delay[sid]
+    buf_offset = sensor_historyadr[sid]
+    t = time_in[worldid] - delay
+    _history_read_vector(
+      0,  # write to result_out starting at index 0 (not global sensor adr)
+      history_in,
+      worldid,
+      buf_offset,
+      nsample,
+      dim,
+      t,
+      interp_val,
+      result_out,
+    )
+
+
+def read_sensor(
+  m: Model,
+  d: Data,
+  sensorid: int,
+  time: wp.array[float],
+  interp: int,
+  result: wp.array2d[float],
+):
+  """Read delayed sensor for 1 sensor across all worlds.
+
+  Args:
+    m: The model containing kinematic and dynamic information.
+    d: The data object containing the current state and output arrays.
+    sensorid: sensor index.
+    time: query time per world (nworld,).
+    interp: interpolation mode (-1=model default, 0=ZOH, 1=linear, 2=cubic).
+    result: output buffer (nworld, dim).
+  """
+  wp.launch(
+    _read_sensor_kernel,
+    dim=(d.nworld,),
+    inputs=[
+      m.sensor_dim,
+      m.sensor_adr,
+      m.sensor_history,
+      m.sensor_historyadr,
+      m.sensor_delay,
+      time,
+      d.history,
+      d.sensordata,
+      sensorid,
+      interp,
+    ],
+    outputs=[result],
+  )
+
+
+@wp.kernel
+def _init_ctrl_history_kernel(
+  # kernel_analyzer: off
+  # Model:
+  actuator_history: wp.array[wp.vec2i],
+  actuator_historyadr: wp.array[int],
+  # In:
+  ctrlid: int,
+  times: wp.array[float],
+  values: wp.array2d[float],
+  has_times: int,
+  # Data out:
+  history_out: wp.array2d[float],
+  # kernel_analyzer: on
+):
+  """Initialize history buffer for 1 actuator across all worlds."""
+  worldid = wp.tid()
+
+  nsample = actuator_history[ctrlid][0]
+  buf_offset = actuator_historyadr[ctrlid]
+
+  # preserve user slot
+  user = history_out[worldid, buf_offset]
+
+  # cursor = 0 (samples in order, newest at index nsample-1)
+  history_out[worldid, buf_offset + 1] = float(nsample - 1)
+
+  times_offset = buf_offset + 2
+  values_offset = buf_offset + 2 + nsample
+
+  for i in range(nsample):
+    if has_times != 0:
+      history_out[worldid, times_offset + i] = times[i]
+    else:
+      history_out[worldid, times_offset + i] = -MJ_MAXVAL
+    history_out[worldid, values_offset + i] = values[worldid, i]
+
+  # restore user slot
+  history_out[worldid, buf_offset] = user
+
+
+def init_ctrl_history(
+  m: Model,
+  d: Data,
+  ctrlid: int,
+  times: wp.array[float],
+  values: wp.array2d[float],
+):
+  """Initialize history buffer for 1 actuator across all worlds.
+
+  Args:
+    m: The model containing kinematic and dynamic information.
+    d: The data object containing the current state and output arrays.
+    ctrlid: actuator index.
+    times: timestamps or None (nsample,).
+    values: ctrl values (nworld, nsample).
+
+  Raises:
+    ValueError: If times are not strictly increasing.
+  """
+  has_times = 0 if times is None else 1
+  if times is not None:
+    t_np = times.numpy()
+    for i in range(len(t_np) - 1):
+      if t_np[i + 1] - t_np[i] < MJ_MINVAL:
+        raise ValueError(f"times must be strictly increasing, got times[{i}]={t_np[i]} >= times[{i + 1}]={t_np[i + 1]}")
+  if times is None:
+    times = wp.empty(0, dtype=float)
+
+  wp.launch(
+    _init_ctrl_history_kernel,
+    dim=(d.nworld,),
+    inputs=[
+      m.actuator_history,
+      m.actuator_historyadr,
+      ctrlid,
+      times,
+      values,
+      has_times,
+    ],
+    outputs=[d.history],
+  )
+
+
+# kernel_analyzer: off
+@wp.kernel
+def _init_sensor_history_kernel(
+  # Model:
+  sensor_history: wp.array[wp.vec2i],
+  sensor_historyadr: wp.array[int],
+  sensor_dim_arr: wp.array[int],
+  # In:
+  sensorid: int,
+  times: wp.array[float],
+  values: wp.array2d[float],
+  phase: wp.array[float],
+  has_times: int,
+  # Data out:
+  history_out: wp.array2d[float],
+):
+  # kernel_analyzer: on
+  """Initialize history buffer for 1 sensor across all worlds."""
+  worldid = wp.tid()
+
+  nsample = sensor_history[sensorid][0]
+  dim = sensor_dim_arr[sensorid]
+  buf_offset = sensor_historyadr[sensorid]
+
+  # set user slot (phase = last computation time for interval sensors)
+  history_out[worldid, buf_offset] = phase[worldid]
+
+  # cursor = 0 (samples in order, newest at index nsample-1)
+  history_out[worldid, buf_offset + 1] = float(nsample - 1)
+
+  times_offset = buf_offset + 2
+  values_offset = buf_offset + 2 + nsample
+
+  for i in range(nsample):
+    if has_times != 0:
+      history_out[worldid, times_offset + i] = times[i]
+    else:
+      history_out[worldid, times_offset + i] = -MJ_MAXVAL
+    for j in range(dim):
+      history_out[worldid, values_offset + i * dim + j] = values[worldid, i * dim + j]
+
+
+def init_sensor_history(
+  m: Model,
+  d: Data,
+  sensorid: int,
+  times: wp.array[float],
+  values: wp.array2d[float],
+  phase: wp.array[float],
+):
+  """Initialize history buffer for 1 sensor across all worlds.
+
+  Args:
+    m: The model containing kinematic and dynamic information.
+    d: The data object containing the current state and output arrays.
+    sensorid: sensor index.
+    times: timestamps or None (nsample,).
+    values: sensor values (nworld, nsample * dim).
+    phase: user slot value per world (nworld,).
+
+  Raises:
+    ValueError: If times are not strictly increasing.
+  """
+  has_times = 0 if times is None else 1
+  if times is not None:
+    t_np = times.numpy()
+    for i in range(len(t_np) - 1):
+      if t_np[i + 1] - t_np[i] < MJ_MINVAL:
+        raise ValueError(f"times must be strictly increasing, got times[{i}]={t_np[i]} >= times[{i + 1}]={t_np[i + 1]}")
+  if times is None:
+    times = wp.empty(0, dtype=float)
+
+  wp.launch(
+    _init_sensor_history_kernel,
+    dim=(d.nworld,),
+    inputs=[
+      m.sensor_history,
+      m.sensor_historyadr,
+      m.sensor_dim,
+      sensorid,
+      times,
+      values,
+      phase,
+      has_times,
+    ],
+    outputs=[d.history],
+  )
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/inverse.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/inverse.py
index 82b7cf4a..388ad1aa 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/inverse.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/inverse.py
@@ -95,9 +95,9 @@ def discrete_acc(m: Model, d: Data, qacc: wp.array2d[float]):
 
     # TODO(team): qacc = d.qacc if (m.dof_damping == 0.0).all()
 
-    # set qfrc = (d.qM + m.opt.timestep * diag(m.dof_damping)) * d.qacc
+    # set qfrc = (d.M + m.opt.timestep * diag(m.dof_damping)) * d.qacc
 
-    # d.qM @ d.qacc
+    # d.M @ d.qacc
     support.mul_m(m, d, qfrc, d.qacc)
 
     # qfrc += m.opt.timestep * damp_deriv * d.qacc
@@ -109,16 +109,16 @@ def discrete_acc(m: Model, d: Data, qacc: wp.array2d[float]):
     )
   elif m.opt.integrator == IntegratorType.IMPLICITFAST:
     if m.is_sparse:
-      qDeriv = wp.empty((d.nworld, 1, m.nM), dtype=float)
+      qDeriv = wp.empty((d.nworld, 1, m.nC), dtype=float)
     else:
       qDeriv = wp.empty((d.nworld, m.nv, m.nv), dtype=float)
     derivative.deriv_smooth_vel(m, d, qDeriv)
     mul_m(m, d, qfrc, d.qacc, M=qDeriv)
-    smooth.factor_solve_i(m, d, d.qM, d.qLD, d.qLDiagInv, qacc, qfrc)
+    smooth.factor_solve_i(m, d, d.M, d.qLD, d.qLDiagInv, qacc, qfrc)
   else:
     raise NotImplementedError(f"integrator {m.opt.integrator} not implemented.")
 
-  # solve for qacc: qfrc = d.qM @ d.qacc
+  # solve for qacc: qfrc = d.M @ d.qacc
   smooth.solve_m(m, d, qacc, qfrc)
 
 
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py
index 309f9e41..ffd9f8f3 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py
@@ -33,6 +33,9 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import TrnType
 from mujoco.mjx.third_party.mujoco_warp._src.types import vec10
 from mujoco.mjx.third_party.mujoco_warp._src.util_pkg import check_version
 
+# TODO(team): remove after improving island solver performance
+ENABLE_ISLANDS = False
+
 
 def _is_array_spec(typ) -> bool:
   """Check if a type annotation is an array spec (wp.array instance or bracket annotation)."""
@@ -64,6 +67,61 @@ def _create_array(data: Any, spec, sizes: dict[str, int]) -> wp.array | None:
   return array
 
 
+def _create_constraint(
+  mjm,
+  nworld: int,
+  njmax: int,
+  njmax_nnz: int,
+  sizes: dict,
+  island_enabled: bool,
+  mjd=None,
+) -> types.Constraint:
+  """Construct a types.Constraint with standard and island local fields allocated properly."""
+  efc_kwargs = {"J_rownnz": None, "J_rowadr": None, "J_colind": None, "J": None}
+  sparse = is_sparse(mjm)
+
+  for f in dataclasses.fields(types.Constraint):
+    if f.name == "itype":
+      efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0), dtype=int)
+    elif f.name == "iid":
+      efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0), dtype=int)
+    elif f.name == "iJ_rownnz":
+      efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0) if sparse else (nworld, 0), dtype=int)
+    elif f.name == "iJ_rowadr":
+      efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0) if sparse else (nworld, 0), dtype=int)
+    elif f.name == "iJ_colind":
+      efc_kwargs[f.name] = wp.empty((nworld, 1, njmax_nnz if island_enabled else 0) if sparse else (nworld, 0, 0), dtype=int)
+    elif f.name == "iJ":
+      efc_kwargs[f.name] = wp.empty(
+        (nworld, 1, njmax_nnz if island_enabled else 0) if sparse else (nworld, njmax if island_enabled else 0, mjm.nv),
+        dtype=float,
+      )
+    elif f.name == "iD":
+      efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0), dtype=float)
+    elif f.name == "iaref":
+      efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0), dtype=float)
+    elif f.name == "ifrictionloss":
+      efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0), dtype=float)
+    elif f.name == "iforce":
+      efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0), dtype=float)
+    elif f.name == "istate":
+      efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0), dtype=int)
+    else:
+      if f.name in efc_kwargs:
+        continue
+
+      if mjd is not None:
+        shape = tuple(sizes[dim] if isinstance(dim, str) else dim for dim in f.type.shape)
+        val = np.zeros(shape, dtype=f.type.dtype)
+        if f.name in ("type", "id", "pos", "margin", "D", "vel", "aref", "frictionloss", "force"):
+          val[:, : mjd.nefc] = np.tile(getattr(mjd, "efc_" + f.name), (nworld, 1))
+        efc_kwargs[f.name] = wp.array(val, dtype=f.type.dtype)
+      else:
+        efc_kwargs[f.name] = _create_array(None, f.type, sizes)
+
+  return types.Constraint(**efc_kwargs)
+
+
 def is_sparse(mjm: mujoco.MjModel) -> bool:
   if mjm.opt.jacobian == mujoco.mjtJacobian.mjJAC_AUTO:
     if mjm.nv > 32:
@@ -183,6 +241,12 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
     opt_kwargs["impratio_invsqrt"] = 1.0 / np.sqrt(np.maximum(mjm.opt.impratio, mujoco.mjMINVAL))
   opt = types.Option(**opt_kwargs)
 
+  # islands are disabled by default while performance is being improved
+  # override by setting io.ENABLE_ISLANDS = True
+  # TODO(team): remove after improving island solver performance
+  if not ENABLE_ISLANDS:
+    opt.disableflags |= types.DisableBit.ISLAND
+
   # C MuJoCo tolerance was chosen for float64 architecture, but we default to float32 on GPU
   # adjust the tolerance for lower precision, to avoid the solver spending iterations needlessly
   # bouncing around the optimal solution
@@ -271,7 +335,7 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
   jnt_limited_slide_hinge = mjm.jnt_limited & np.isin(mjm.jnt_type, (mujoco.mjtJoint.mjJNT_SLIDE, mujoco.mjtJoint.mjJNT_HINGE))
   m.jnt_limited_slide_hinge_adr = np.nonzero(jnt_limited_slide_hinge)[0]
   m.jnt_limited_ball_adr = np.nonzero(mjm.jnt_limited & (mjm.jnt_type == mujoco.mjtJoint.mjJNT_BALL))[0]
-  m.dof_tri_row, m.dof_tri_col = np.tril_indices(mjm.nv)
+  m.dof_tri_row, m.dof_tri_col = np.triu_indices(mjm.nv)
 
   # precompute body_isdofancestor: which DOFs affect each body
   # TODO: Investigate alternative approach such as bitmap
@@ -585,14 +649,14 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
     for j in range(mjm.mesh_vertnum[mjm.sensor_objid[i]])
   ]
 
-  # qM_tiles records the block diagonal structure of qM
+  # M_tiles records the block diagonal structure of M
   tile_corners = [i for i in range(mjm.nv) if mjm.dof_parentid[i] == -1]
   tiles = {}
   for i in range(len(tile_corners)):
     tile_beg = tile_corners[i]
     tile_end = mjm.nv if i == len(tile_corners) - 1 else tile_corners[i + 1]
     tiles.setdefault(tile_end - tile_beg, []).append(tile_beg)
-  m.qM_tiles = tuple(types.TileSet(adr=wp.array(tiles[sz], dtype=int), size=sz) for sz in sorted(tiles.keys()))
+  m.M_tiles = tuple(types.TileSet(adr=wp.array(tiles[sz], dtype=int), size=sz) for sz in sorted(tiles.keys()))
 
   # qLD_updates has dof tree ordering of qLD updates for sparse factor m
   qLD_updates, dof_depth = {}, np.zeros(mjm.nv, dtype=int) - 1
@@ -620,41 +684,69 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
   m.qLD_all_updates = all_updates_flat if all_updates_flat else [(0, 0, 0)]
   m.qLD_level_offsets = level_offsets
 
-  # indices for sparse qM_fullm (used in solver)
-  m.qM_fullm_i, m.qM_fullm_j = [], []
+  # Indices for sparse M_fullm (used in solver). M_fullm_i/j are built by
+  # walking dof_parentid for each dof, so for joint types whose internal block
+  # MuJoCo stores diagonal-only in the compact (M_rownnz, M_rowadr) layout
+  # (e.g. free joints), the chain-aware layout here has more entries per row
+  # than the compact layout.
+  m.M_fullm_i, m.M_fullm_j = [], []
   for i in range(mjm.nv):
     j = i
     while j > -1:
-      m.qM_fullm_i.append(i)
-      m.qM_fullm_j.append(j)
+      m.M_fullm_i.append(i)
+      m.M_fullm_j.append(j)
       j = mjm.dof_parentid[j]
+  # M_elemid maps (row, col) -> madr index in the native CSR M layout
+  M_elemid = np.full((mjm.nv, mjm.nv), -1, dtype=np.int32)
+  for i in range(mjm.nv):
+    rowadr = mjm.M_rowadr[i]
+    rownnz = mjm.M_rownnz[i]
+    for k in range(rownnz):
+      madr = rowadr + k
+      col = int(mjm.M_colind[madr])
+      M_elemid[i, col] = madr
+  m.M_elemid = M_elemid
+
+  upper_j, upper_i = np.triu_indices(mjm.nv)
+  upper_elemid = M_elemid[upper_i, upper_j]
+  valid_mask = upper_elemid != -1
+  m.M_fullm_upper_i = upper_j[valid_mask].tolist()
+  m.M_fullm_upper_j = upper_i[valid_mask].tolist()
+  m.M_fullm_upper_elemid = upper_elemid[valid_mask].tolist()
+
+  # indices for sparse qD_fullm (used in RNE derivatives)
+  # D-structure is the full square sparsity pattern (both upper and lower triangle)
+  m.qD_fullm_i, m.qD_fullm_j = [], []
+  for i in range(mjm.nv):
+    rowadr = mjm.D_rowadr[i]
+    rownnz = mjm.D_rownnz[i]
+    for k in range(rownnz):
+      m.qD_fullm_i.append(i)
+      m.qD_fullm_j.append(int(mjm.D_colind[rowadr + k]))
+  m.nD = mjm.nD
 
   # Gather-based sparse mul_m: for each row, all (col, madr) including diagonal
   row_elements = [[] for _ in range(mjm.nv)]
 
-  # Add diagonal
   for i in range(mjm.nv):
-    row_elements[i].append((i, mjm.dof_Madr[i]))
-
-  # Add off-diagonals: ancestors (lower) and descendants (upper)
-  for i in range(mjm.nv):
-    madr_ij, j = mjm.dof_Madr[i], i
-    while True:
-      madr_ij, j = madr_ij + 1, mjm.dof_parentid[j]
-      if j == -1:
-        break
-      row_elements[i].append((j, madr_ij))  # row i gathers M[i,j] * vec[j]
-      row_elements[j].append((i, madr_ij))  # row j gathers M[j,i] * vec[i]
+    rowadr = mjm.M_rowadr[i]
+    rownnz = mjm.M_rownnz[i]
+    for k in range(rownnz):
+      madr = rowadr + k
+      col = int(mjm.M_colind[madr])
+      row_elements[i].append((col, madr))  # row i gathers M[i,col] * vec[col]
+      if i != col:
+        row_elements[col].append((i, madr))  # row col gathers M[i,col] * vec[i]
 
   # Flatten into CSR-like arrays
-  m.qM_mulm_rowadr = [0]
-  m.qM_mulm_col = []
-  m.qM_mulm_madr = []
+  m.M_mulm_rowadr = [0]
+  m.M_mulm_col = []
+  m.M_mulm_madr = []
   for i in range(mjm.nv):
     for col, madr in row_elements[i]:
-      m.qM_mulm_col.append(col)
-      m.qM_mulm_madr.append(madr)
-    m.qM_mulm_rowadr.append(len(m.qM_mulm_col))
+      m.M_mulm_col.append(col)
+      m.M_mulm_madr.append(madr)
+    m.M_mulm_rowadr.append(len(m.M_mulm_col))
 
   m.flexedge_J_rownnz = mjm.flexedge_J_rownnz
   m.flexedge_J_rowadr = mjm.flexedge_J_rowadr
@@ -886,6 +978,42 @@ def _resolve_batch_size(na: int | None, n: int | None, nworld: int, default: int
   return default
 
 
+def _allocate_island_arrays(
+  mjm: mujoco.MjModel,
+  d: types.Data,
+  nworld: int,
+  njmax: int,
+  island_enabled: bool,
+  mjd: mujoco.MjData,
+):
+  ntree_size = mjm.ntree if island_enabled else 0
+  nv_size = mjm.nv if island_enabled else 0
+  njmax_size = njmax if island_enabled else 0
+
+  d.nisland = wp.array(np.full(nworld, mjd.nisland), dtype=int)
+  d.tree_island = wp.array(np.tile(mjd.tree_island, (nworld, 1 if island_enabled else 0)), dtype=int)
+  d.dof_island = wp.array(np.tile(mjd.dof_island, (nworld, 1 if island_enabled else 0)), dtype=int)
+
+  d.island_dofadr = wp.empty((nworld, ntree_size), dtype=int)
+  d.island_nv = wp.empty((nworld, ntree_size), dtype=int)
+  d.island_nefc = wp.empty((nworld, ntree_size), dtype=int)
+  d.island_ne = wp.empty((nworld, ntree_size), dtype=int)
+  d.island_nf = wp.empty((nworld, ntree_size), dtype=int)
+  d.island_efcadr = wp.empty((nworld, ntree_size), dtype=int)
+  d.nidof = wp.empty((nworld if island_enabled else 0,), dtype=int)
+  d.map_dof2idof = wp.empty((nworld, nv_size), dtype=int)
+  d.map_idof2dof = wp.empty((nworld, nv_size), dtype=int)
+  d.map_efc2iefc = wp.empty((nworld, njmax_size), dtype=int)
+  d.map_iefc2efc = wp.empty((nworld, njmax_size), dtype=int)
+
+  d.dof_islandid = wp.empty((nworld, nv_size), dtype=int)
+  d.efc_islandid = wp.empty((nworld, njmax_size), dtype=int)
+  d.iqacc = wp.empty((nworld, nv_size), dtype=float)
+  d.iqacc_smooth = wp.empty((nworld, nv_size), dtype=float)
+  d.iqfrc_smooth = wp.empty((nworld, nv_size), dtype=float)
+  d.iqfrc_constraint = wp.empty((nworld, nv_size), dtype=float)
+
+
 def make_data(
   mjm: mujoco.MjModel,
   nworld: int = 1,
@@ -967,7 +1095,8 @@ def make_data(
 
   contact = types.Contact(**{f.name: _create_array(None, f.type, sizes) for f in dataclasses.fields(types.Contact)})
   contact.efc_address = wp.array(np.full((naconmax, sizes["nmaxpyramid"]), -1, dtype=int), dtype=int)
-  efc = types.Constraint(**{f.name: _create_array(None, f.type, sizes) for f in dataclasses.fields(types.Constraint)})
+
+  efc = _create_constraint(mjm, nworld, njmax, njmax_nnz, sizes, ENABLE_ISLANDS)
 
   if is_sparse(mjm):
     efc.J_rownnz = wp.zeros((nworld, njmax), dtype=int)
@@ -1005,7 +1134,7 @@ def make_data(
     "njmax": njmax,
     "njmax_pad": sizes["njmax_pad"],
     "njmax_nnz": njmax_nnz,
-    "qM": None,
+    "M": None,
     "qLD": None,
     # world body
     "xquat": wp.array(np.tile(mjd.xquat, (nworld, 1)), shape=(nworld, mjm.nbody), dtype=wp.quat),
@@ -1024,6 +1153,24 @@ def make_data(
     # island arrays
     "nisland": None,
     "tree_island": None,
+    "dof_island": None,
+    "island_dofadr": None,
+    "island_nv": None,
+    "island_nefc": None,
+    "island_ne": None,
+    "island_nf": None,
+    "island_efcadr": None,
+    "nidof": None,
+    "map_dof2idof": None,
+    "map_idof2dof": None,
+    "map_efc2iefc": None,
+    "map_iefc2efc": None,
+    "dof_islandid": None,
+    "efc_islandid": None,
+    "iqacc": None,
+    "iqacc_smooth": None,
+    "iqfrc_smooth": None,
+    "iqfrc_constraint": None,
   }
   for f in dataclasses.fields(types.Data):
     if f.name in d_kwargs:
@@ -1033,15 +1180,13 @@ def make_data(
   d = types.Data(**d_kwargs)
 
   if is_sparse(mjm):
-    d.qM = wp.zeros((nworld, 1, mjm.nM), dtype=float)
+    d.M = wp.zeros((nworld, 1, mjm.nC), dtype=float)
     d.qLD = wp.zeros((nworld, 1, mjm.nC), dtype=float)
   else:
-    d.qM = wp.zeros((nworld, sizes["nv_pad"], sizes["nv_pad"]), dtype=float)
+    d.M = wp.zeros((nworld, sizes["nv_pad"], sizes["nv_pad"]), dtype=float)
     d.qLD = wp.zeros((nworld, mjm.nv, mjm.nv), dtype=float)
 
-  # island discovery arrays
-  d.nisland = wp.zeros((nworld,), dtype=int)
-  d.tree_island = wp.zeros((nworld, mjm.ntree), dtype=int)
+  _allocate_island_arrays(mjm, d, nworld, njmax, ENABLE_ISLANDS, mjd)
 
   return d
 
@@ -1176,16 +1321,7 @@ def put_data(
   # create efc
   efc_kwargs = {"J_rownnz": None, "J_rowadr": None, "J_colind": None, "J": None}
 
-  for f in dataclasses.fields(types.Constraint):
-    if f.name in efc_kwargs:
-      continue
-    shape = tuple(sizes[dim] if isinstance(dim, str) else dim for dim in f.type.shape)
-    val = np.zeros(shape, dtype=f.type.dtype)
-    if f.name in ("type", "id", "pos", "margin", "D", "vel", "aref", "frictionloss", "force"):
-      val[:, : mjd.nefc] = np.tile(getattr(mjd, "efc_" + f.name), (nworld, 1))
-    efc_kwargs[f.name] = wp.array(val, dtype=f.type.dtype)
-
-  efc = types.Constraint(**efc_kwargs)
+  efc = _create_constraint(mjm, nworld, njmax, njmax_nnz, sizes, ENABLE_ISLANDS, mjd)
 
   if is_sparse(mjm):
     J_rownnz = np.zeros(njmax, dtype=np.int32)
@@ -1236,12 +1372,30 @@ def put_data(
     "njmax_nnz": njmax_nnz,
     # fields set after initialization:
     "solver_niter": None,
-    "qM": None,
+    "M": None,
     "qLD": None,
     "nacon": None,
     # island arrays
     "nisland": None,
     "tree_island": None,
+    "dof_island": None,
+    "island_dofadr": None,
+    "island_nv": None,
+    "island_nefc": None,
+    "island_ne": None,
+    "island_nf": None,
+    "island_efcadr": None,
+    "nidof": None,
+    "map_dof2idof": None,
+    "map_idof2dof": None,
+    "map_efc2iefc": None,
+    "map_iefc2efc": None,
+    "dof_islandid": None,
+    "efc_islandid": None,
+    "iqacc": None,
+    "iqacc_smooth": None,
+    "iqfrc_smooth": None,
+    "iqfrc_constraint": None,
   }
   for f in dataclasses.fields(types.Data):
     if f.name in d_kwargs:
@@ -1256,20 +1410,25 @@ def put_data(
   d.solver_niter = wp.full((nworld,), mjd.solver_niter[0], dtype=int)
 
   if is_sparse(mjm):
-    d.qM = wp.array(np.full((nworld, 1, mjm.nM), mjd.qM), dtype=float)
+    if check_version("mujoco>=3.8.1.dev910242375"):
+      d.M = wp.array(np.full((nworld, 1, mjm.nC), mjd.M), dtype=float)
+    else:
+      d.M = wp.array(np.full((nworld, 1, mjm.nC), mjd.qM[mjm.mapM2M]), dtype=float)
     d.qLD = wp.array(np.full((nworld, 1, mjm.nC), mjd.qLD), dtype=float)
   else:
-    qM = np.zeros((mjm.nv, mjm.nv))
-    mujoco.mj_fullM(mjm, qM, mjd.qM)
-    qLD = np.linalg.cholesky(qM) if (mjd.qM != 0.0).any() and (mjd.qLD != 0.0).any() else np.zeros((mjm.nv, mjm.nv))
+    M = np.zeros((mjm.nv, mjm.nv))
+    if check_version("mujoco>=3.8.1.dev910242375"):
+      mujoco.mju_sym2dense(M, mjd.M, mjm.M_rownnz, mjm.M_rowadr, mjm.M_colind)
+      qLD = np.linalg.cholesky(M).T if (mjd.M != 0.0).any() and (mjd.qLD != 0.0).any() else np.zeros((mjm.nv, mjm.nv))
+    else:
+      mujoco.mj_fullM(mjm, M, mjd.qM)
+      qLD = np.linalg.cholesky(M).T if (mjd.qM != 0.0).any() and (mjd.qLD != 0.0).any() else np.zeros((mjm.nv, mjm.nv))
     padding = sizes["nv_pad"] - mjm.nv
-    qM_padded = np.pad(qM, ((0, padding), (0, padding)), mode="constant", constant_values=0.0)
-    d.qM = wp.array(np.full((nworld, sizes["nv_pad"], sizes["nv_pad"]), qM_padded), dtype=float)
+    M_padded = np.pad(M, ((0, padding), (0, padding)), mode="constant", constant_values=0.0)
+    d.M = wp.array(np.full((nworld, sizes["nv_pad"], sizes["nv_pad"]), M_padded), dtype=float)
     d.qLD = wp.array(np.full((nworld, mjm.nv, mjm.nv), qLD), dtype=float)
 
-  # island arrays
-  d.nisland = wp.array(np.full(nworld, mjd.nisland), dtype=int)
-  d.tree_island = wp.array(np.tile(mjd.tree_island, (nworld, 1)), dtype=int)
+  _allocate_island_arrays(mjm, d, nworld, njmax, ENABLE_ISLANDS, mjd)
 
   d.nacon = wp.array([mjd.ncon * nworld], dtype=int)
 
@@ -1403,6 +1562,9 @@ def get_data_into(
   result.qfrc_constraint[:] = d.qfrc_constraint.numpy()[world_id]
   result.qfrc_inverse[:] = d.qfrc_inverse.numpy()[world_id]
 
+  if mjm.nhistory > 0:
+    result.history[:] = d.history.numpy()[world_id]
+
   # contact
   result.contact.dist[:ncon] = d.contact.dist.numpy()[ncon_filter]
   result.contact.pos[:ncon] = d.contact.pos.numpy()[ncon_filter]
@@ -1417,17 +1579,27 @@ def get_data_into(
   result.contact.efc_address[:ncon] = contact_efc_address_ordered[:ncon]
 
   if is_sparse(mjm):
-    result.qM[:] = d.qM.numpy()[world_id, 0]
+    if check_version("mujoco>=3.8.1.dev910242375"):
+      result.M[:] = d.M.numpy()[world_id, 0]
+    else:
+      result.qM[mjm.mapM2M] = d.M.numpy()[world_id, 0]
     result.qLD[:] = d.qLD.numpy()[world_id, 0]
   else:
-    qM = d.qM.numpy()[world_id]
-    adr = 0
-    for i in range(mjm.nv):
-      j = i
-      while j >= 0:
-        result.qM[adr] = qM[i, j]
-        j = mjm.dof_parentid[j]
-        adr += 1
+    M = d.M.numpy()[world_id]
+    if check_version("mujoco>=3.8.1.dev910242375"):
+      for i in range(mjm.nv):
+        adr = mjm.M_rowadr[i]
+        for k in range(mjm.M_rownnz[i]):
+          col = mjm.M_colind[adr + k]
+          result.M[adr + k] = M[i, col]
+    else:
+      adr = 0
+      for i in range(mjm.nv):
+        j = i
+        while j >= 0:
+          result.qM[adr] = M[i, j]
+          j = mjm.dof_parentid[j]
+          adr += 1
     mujoco.mj_factorM(mjm, result)
 
   if nefc > 0:
@@ -1466,6 +1638,7 @@ def get_data_into(
   result.efc_frictionloss[:] = d.efc.frictionloss.numpy()[world_id, efc_idx]
   result.efc_state[:] = d.efc.state.numpy()[world_id, efc_idx]
   result.efc_force[:] = d.efc.force.numpy()[world_id, efc_idx]
+  result.efc_island[:] = d.efc.island.numpy()[world_id, efc_idx]
 
   # rne_postconstraint
   result.cacc[:] = d.cacc.numpy()[world_id]
@@ -1487,8 +1660,51 @@ def get_data_into(
   # islands
   nisland = d.nisland.numpy()[world_id]
   result.nisland = nisland
-  if nisland:
+  if d.tree_island.shape[1] > 0 and nisland:
     result.tree_island[:] = d.tree_island.numpy()[world_id]
+    result.dof_island[:] = d.dof_island.numpy()[world_id]
+    result.island_dofadr[:nisland] = d.island_dofadr.numpy()[world_id, :nisland]
+    result.island_nv[:nisland] = d.island_nv.numpy()[world_id, :nisland]
+    result.island_nefc[:nisland] = d.island_nefc.numpy()[world_id, :nisland]
+    result.island_ne[:nisland] = d.island_ne.numpy()[world_id, :nisland]
+    result.island_nf[:nisland] = d.island_nf.numpy()[world_id, :nisland]
+    result.island_iefcadr[:nisland] = d.island_efcadr.numpy()[world_id, :nisland]
+    nv = mjm.nv
+    result.map_dof2idof[:nv] = d.map_dof2idof.numpy()[world_id, :nv]
+    result.map_idof2dof[:nv] = d.map_idof2dof.numpy()[world_id, :nv]
+    result.map_efc2iefc[:nefc] = d.map_efc2iefc.numpy()[world_id, :nefc]
+    result.map_iefc2efc[:nefc] = d.map_iefc2efc.numpy()[world_id, :nefc]
+
+    result.iefc_type[:nefc] = d.efc.itype.numpy()[world_id, :nefc]
+    result.iefc_id[:nefc] = d.efc.iid.numpy()[world_id, :nefc]
+    result.iefc_D[:nefc] = d.efc.iD.numpy()[world_id, :nefc]
+    result.iefc_aref[:nefc] = d.efc.iaref.numpy()[world_id, :nefc]
+    result.iefc_frictionloss[:nefc] = d.efc.ifrictionloss.numpy()[world_id, :nefc]
+    result.iefc_state[:nefc] = d.efc.istate.numpy()[world_id, :nefc]
+    result.iefc_force[:nefc] = d.efc.iforce.numpy()[world_id, :nefc]
+
+    if is_sparse(mjm):
+      iefc_J = np.zeros((nefc, mjm.nv))
+      mujoco.mju_sparse2dense(
+        iefc_J,
+        d.efc.iJ.numpy()[world_id, 0],
+        d.efc.iJ_rownnz.numpy()[world_id, :nefc],
+        d.efc.iJ_rowadr.numpy()[world_id, :nefc],
+        d.efc.iJ_colind.numpy()[world_id, 0],
+      )
+    else:
+      iefc_J = d.efc.iJ.numpy()[world_id, :nefc, : mjm.nv]
+
+    if mujoco.mj_isSparse(mjm):
+      mujoco.mju_dense2sparse(
+        result.iefc_J,
+        iefc_J,
+        result.iefc_J_rownnz,
+        result.iefc_J_rowadr,
+        result.iefc_J_colind,
+      )
+    else:
+      result.iefc_J[: nefc * mjm.nv] = iefc_J.flatten()
 
 
 def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
@@ -1511,14 +1727,14 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
     xfrc_applied_out[worldid, bodyid][elemid] = 0.0
 
   @wp.kernel(module="unique", enable_backward=False)
-  def reset_qM(reset_in: wp.array[bool], qM_out: wp.array3d[float]):
+  def reset_M(reset_in: wp.array[bool], M_out: wp.array3d[float]):
     worldid, elemid1, elemid2 = wp.tid()
 
     if wp.static(reset is not None):
       if not reset_in[worldid]:
         return
 
-    qM_out[worldid, elemid1, elemid2] = 0.0
+    M_out[worldid, elemid1, elemid2] = 0.0
 
   @wp.kernel(module="unique", enable_backward=False)
   def reset_nworld(
@@ -1670,10 +1886,10 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
 
   wp.launch(reset_xfrc_applied, dim=(d.nworld, m.nbody, 6), inputs=[reset_input], outputs=[d.xfrc_applied])
   wp.launch(
-    reset_qM,
-    dim=(d.nworld, d.qM.shape[1], d.qM.shape[2]),
+    reset_M,
+    dim=(d.nworld, d.M.shape[1], d.M.shape[2]),
     inputs=[reset_input],
-    outputs=[d.qM],
+    outputs=[d.M],
   )
 
   # set mocap_pos/quat = body_pos/quat for mocap bodies
@@ -1786,11 +2002,12 @@ def _copy_tendon_length0(
 def _compute_meaninertia(
   nv: int,
   is_sparse: bool,
-  dof_Madr_in: wp.array[int],
-  qM_in: wp.array3d[float],
+  M_rownnz_in: wp.array[int],
+  M_rowadr_in: wp.array[int],
+  M_in: wp.array3d[float],
   meaninertia_out: wp.array[float],
 ):
-  """Compute mean diagonal inertia from qM at qpos0."""
+  """Compute mean diagonal inertia from M at qpos0."""
   worldid = wp.tid()
 
   if nv == 0:
@@ -1800,12 +2017,12 @@ def _compute_meaninertia(
   total = float(0.0)
   for i in range(nv):
     if is_sparse:
-      # Sparse: qM is flattened lower triangular, diagonal at dof_Madr[i]
-      madr = dof_Madr_in[i]
-      total += qM_in[worldid, 0, madr]
+      # Sparse: M is in CSR format, diagonal at M_rowadr_in[i] + M_rownnz_in[i] - 1
+      madr = M_rowadr_in[i] + M_rownnz_in[i] - 1
+      total += M_in[worldid, 0, madr]
     else:
-      # Dense: qM is 2D matrix, diagonal at [i,i]
-      total += qM_in[worldid, i, i]
+      # Dense: M is 2D matrix, diagonal at [i,i]
+      total += M_in[worldid, i, i]
 
   meaninertia_out[worldid % meaninertia_out.shape[0]] = total / float(nv)
 
@@ -2288,11 +2505,11 @@ def set_const_0(m: types.Model, d: types.Data):
   smooth.factor_m(m, d)
   smooth.transmission(m, d)
 
-  # Compute meaninertia from qM diagonal at qpos0
+  # Compute meaninertia from M diagonal at qpos0
   wp.launch(
     _compute_meaninertia,
     dim=d.nworld,
-    inputs=[m.nv, m.is_sparse, m.dof_Madr, d.qM],
+    inputs=[m.nv, m.is_sparse, m.M_rownnz, m.M_rowadr, d.M],
     outputs=[m.stat.meaninertia],
   )
 
@@ -2452,27 +2669,31 @@ def set_const(m: types.Model, d: types.Data):
 
   Model fields that can be modified safely with set_const:
 
-    Field                            | Notes
-    ---------------------------------|----------------------------------------------
-    qpos0, qpos_spring               |
-    body_mass, body_inertia,         | Mass and inertia are usually scaled together
-    body_ipos, body_iquat            | since inertia is sum(m * r^2).
-    body_pos, body_quat              | Unsafe for static bodies (invalidates BVH).
-    body_gravcomp                    | If changing from 0 to >0 bodies, required.
-    dof_armature                     |
-    eq_data                          | For connect/weld, offsets computed if not set.
-    hfield_size                      |
-    tendon_stiffness, tendon_damping | Only if changing from/to zero.
-    actuator_gainprm, actuator_biasprm | For position actuators with dampratio.
+  ==================================  ==============================================
+  Field                               Notes
+  ==================================  ==============================================
+  qpos0, qpos_spring
+  body_mass, body_inertia,            Mass and inertia are usually scaled together
+  body_ipos, body_iquat               since inertia is sum(m * r^2).
+  body_pos, body_quat                 Unsafe for static bodies (invalidates BVH).
+  body_gravcomp                       If changing from 0 to >0 bodies, required.
+  dof_armature
+  eq_data                             For connect/weld, offsets computed if not set.
+  hfield_size
+  tendon_stiffness, tendon_damping    Only if changing from/to zero.
+  actuator_gainprm, actuator_biasprm  For position actuators with dampratio.
+  ==================================  ==============================================
 
   For selective updates, use the sub-functions directly based on what changed:
 
-    Modified Field  | Call
-    ----------------|------------------
-    body_mass       | set_const
-    body_gravcomp   | set_const_fixed
-    body_inertia    | set_const_0
-    qpos0           | set_const_0
+  ==============  ===============
+  Modified Field  Call
+  ==============  ===============
+  body_mass       set_const
+  body_gravcomp   set_const_fixed
+  body_inertia    set_const_0
+  qpos0           set_const_0
+  ==============  ===============
 
   Computes:
     - Fixed quantities (via set_const_fixed):
@@ -2621,6 +2842,9 @@ def override_model(model: types.Model | mujoco.MjModel, overrides: dict[str, Any
       else:
         val = typ(val)
 
+      if attr == "disableflags" and isinstance(obj, types.Option) and not ENABLE_ISLANDS:
+        val = int(val) | types.DisableBit.ISLAND
+
       setattr(obj, attr, val)
 
 
@@ -2724,11 +2948,13 @@ def create_render_context(
   render_seg: list[bool] | bool | None = None,
   use_textures: bool = True,
   use_shadows: bool = False,
+  use_ambient_lighting: bool = True,
   enabled_geom_groups: list[int] = [0, 1, 2],
   cam_active: list[bool] | None = None,
   flex_render_smooth: bool = True,
   use_precomputed_rays: bool = True,
   render_skybox: bool = False,
+  enable_backface_culling: bool = True,
 ) -> types.RenderContext:
   """Creates a render context on device.
 
@@ -2743,6 +2969,8 @@ def create_render_context(
       If None, uses the MuJoCo model values.
     use_textures: Whether to use textures.
     use_shadows: Whether to use shadows.
+    use_ambient_lighting: Whether to add the renderer's hemispheric ambient
+      lighting term before applying model lights.
     enabled_geom_groups: The geom groups to render.
     cam_active: List of booleans indicating which cameras to include in rendering.
                 If None, all cameras are included.
@@ -2751,6 +2979,10 @@ def create_render_context(
                           When using domain randomization for camera intrinsics, set to False.
     render_skybox: Whether to shade missed rays with the MuJoCo skybox texture.
                    Requires the model to contain a texture with type `mjTEXTURE_SKYBOX`.
+    enable_backface_culling: Drop primitive-ray hits whose normal faces away from
+                             the ray (ray origin inside the geom). Matches MuJoCo's
+                             mesh-ray rule. Default True. Disable for a small
+                             performance gain when no camera is ever inside a geom.
 
   Returns:
     The render context containing rendering fields and output arrays on device.
@@ -2758,13 +2990,7 @@ def create_render_context(
   mjd = mujoco.MjData(mjm)
   mujoco.mj_forward(mjm, mjd)
 
-  constructor = "sah"
-  if check_version("warp>=1.13.0.dev20260325"):
-    # TODO: The cubql constructor and is_cubql_available exist only in
-    # recent Warp 1.13+ builds, modify this after warp is updated to 1.13+.
-    _cubql_avail = getattr(wp, "is_cubql_available", None)
-    if callable(_cubql_avail) and _cubql_avail():
-      constructor = "cubql"
+  constructor = "cubql"
 
   # Mesh BVHs – build for all meshes so per-world variants are available
   nmesh = mjm.nmesh
@@ -2937,6 +3163,7 @@ def create_render_context(
     cam_id_map=wp.array(active_cam_indices, dtype=int),
     use_textures=use_textures,
     use_shadows=use_shadows,
+    use_ambient_lighting=use_ambient_lighting,
     background_color=render_util.pack_rgba_to_uint32(0.1 * 255.0, 0.1 * 255.0, 0.2 * 255.0, 1.0 * 255.0),
     use_precomputed_rays=use_precomputed_rays,
     render_skybox=render_skybox,
@@ -2982,6 +3209,7 @@ def create_render_context(
     render_seg=wp.array(render_seg, dtype=bool),
     znear=znear,
     total_rays=int(total),
+    enable_backface_culling=enable_backface_culling,
   )
 
   bvh.build_scene_bvh(mjm, mjd, rc, nworld)
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/island.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/island.py
index f7fb3a10..0e92b034 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/island.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/island.py
@@ -18,6 +18,7 @@ import warp as wp
 from mujoco.mjx.third_party.mujoco_warp._src import types
 from mujoco.mjx.third_party.mujoco_warp._src.types import ConstraintType
 from mujoco.mjx.third_party.mujoco_warp._src.types import EqType
+from mujoco.mjx.third_party.mujoco_warp._src.types import IslandSolverContext
 from mujoco.mjx.third_party.mujoco_warp._src.types import ObjType
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
 
@@ -35,11 +36,15 @@ def _tree_edges(
   eq_obj1id: wp.array[int],
   eq_obj2id: wp.array[int],
   eq_objtype: wp.array[int],
+  is_sparse: bool,
   # Data in:
   nefc_in: wp.array[int],
   contact_geom_in: wp.array[wp.vec2i],
   efc_type_in: wp.array2d[int],
   efc_id_in: wp.array2d[int],
+  efc_J_rownnz_in: wp.array2d[int],
+  efc_J_rowadr_in: wp.array2d[int],
+  efc_J_colind_in: wp.array3d[int],
   efc_J_in: wp.array3d[float],
   njmax_in: int,
   # Out:
@@ -131,26 +136,39 @@ def _tree_edges(
   first_tree = int(-1)
   has_cross_edge = int(0)
 
-  for dof in range(nv):
-    # TODO(team): sparse efc_J
-    # TODO(team): tree dof skip
-    J_val = efc_J_in[worldid, efcid, dof]
-    if J_val != 0.0:
-      tree = dof_treeid[dof]
-      if tree < 0:
+  count = nv
+  rowadr = 0
+  if is_sparse:
+    count = efc_J_rownnz_in[worldid, efcid]
+    rowadr = efc_J_rowadr_in[worldid, efcid]
+
+  for i in range(count):
+    dof = i
+    if is_sparse:
+      sparseid = rowadr + i
+      dof = efc_J_colind_in[worldid, 0, sparseid]
+    else:
+      J_val = efc_J_in[worldid, efcid, dof]
+      if J_val == 0.0:
         continue
-      if first_tree == -1:
-        first_tree = tree
-      elif tree != first_tree:
-        t1 = wp.min(first_tree, tree)
-        t2 = wp.max(first_tree, tree)
-        wp.atomic_max(tree_tree, worldid, t1, t2, 1)
-        has_cross_edge = 1
+
+    tree = dof_treeid[dof]
+    if tree < 0:
+      continue
+    if first_tree == -1:
+      first_tree = tree
+    elif tree != first_tree:
+      t1 = wp.min(first_tree, tree)
+      t2 = wp.max(first_tree, tree)
+      wp.atomic_max(tree_tree, worldid, t1, t2, 1)
+      wp.atomic_max(tree_tree, worldid, t2, t1, 1)
+      has_cross_edge = 1
 
   if first_tree >= 0 and has_cross_edge == 0:
     wp.atomic_max(tree_tree, worldid, first_tree, first_tree, 1)
 
 
+@event_scope
 def tree_edges(m: types.Model, d: types.Data, tree_tree: wp.array3d[int]):
   """Compute tree-tree adjacency matrix."""
   tree_tree.zero_()
@@ -168,10 +186,14 @@ def tree_edges(m: types.Model, d: types.Data, tree_tree: wp.array3d[int]):
       m.eq_obj1id,
       m.eq_obj2id,
       m.eq_objtype,
+      m.is_sparse,
       d.nefc,
       d.contact.geom,
       d.efc.type,
       d.efc.id,
+      d.efc.J_rownnz,
+      d.efc.J_rowadr,
+      d.efc.J_colind,
       d.efc.J,
       d.njmax,
     ],
@@ -242,6 +264,19 @@ def _flood_fill(
   nisland_out[worldid] = nisland
 
 
+@event_scope
+def flood_fill(m: types.Model, d: types.Data, tree_tree: wp.array3d[int]):
+  d.tree_island.fill_(-1)
+  stack_scratch = wp.empty((d.nworld, m.ntree * m.ntree), dtype=int)
+
+  wp.launch(
+    _flood_fill,
+    dim=d.nworld,
+    inputs=[m.ntree, tree_tree, d.tree_island, stack_scratch],
+    outputs=[d.nisland, d.tree_island, stack_scratch],
+  )
+
+
 @event_scope
 def island(m: types.Model, d: types.Data):
   """Discover constraint islands."""
@@ -254,12 +289,778 @@ def island(m: types.Model, d: types.Data):
   tree_edges(m, d, tree_tree)
 
   # Step 2: DFS flood fill
-  d.tree_island.fill_(-1)
-  stack_scratch = wp.empty((d.nworld, m.ntree * m.ntree), dtype=int)
+  flood_fill(m, d, tree_tree)
+
+
+@wp.kernel
+def _island_count_dofs(
+  dof_treeid: wp.array[int],
+  tree_island_in: wp.array2d[int],
+  dof_island_out: wp.array2d[int],
+  island_nv_out: wp.array2d[int],
+):
+  worldid, dofid = wp.tid()
+
+  island_id = tree_island_in[worldid, dof_treeid[dofid]]
+  dof_island_out[worldid, dofid] = island_id
+  if island_id >= 0:
+    wp.atomic_add(island_nv_out, worldid, island_id, 1)
+
+
+@wp.kernel
+def _island_scan_sizes(
+  nisland_in: wp.array[int],
+  island_idofadr_out: wp.array2d[int],
+  island_nv_inout: wp.array2d[int],
+  island_nefc_inout: wp.array2d[int],
+  island_iefcadr_out: wp.array2d[int],
+  nidof_out: wp.array[int],
+):
+  worldid = wp.tid()
+
+  nisland = nisland_in[worldid]
+  if nisland == 0:
+    nidof_out[worldid] = 0
+    return
+
+  # Scan DOFs and Constraints
+  island_idofadr_out[worldid, 0] = 0
+  island_iefcadr_out[worldid, 0] = 0
+  for i in range(1, nisland):
+    island_idofadr_out[worldid, i] = island_idofadr_out[worldid, i - 1] + island_nv_inout[worldid, i - 1]
+    island_iefcadr_out[worldid, i] = island_iefcadr_out[worldid, i - 1] + island_nefc_inout[worldid, i - 1]
+
+  nidof = island_idofadr_out[worldid, nisland - 1] + island_nv_inout[worldid, nisland - 1]
+  nidof_out[worldid] = nidof
+
+  # Reset for recount
+  for i in range(nisland):
+    island_nv_inout[worldid, i] = 0
+    island_nefc_inout[worldid, i] = 0
+
+
+@wp.kernel
+def _island_map_dofs(
+  nv: int,
+  dof_island_in: wp.array2d[int],
+  island_idofadr_in: wp.array2d[int],
+  nidof_in: wp.array[int],
+  island_nv_inout: wp.array2d[int],
+  map_dof2idof_out: wp.array2d[int],
+  map_idof2dof_out: wp.array2d[int],
+  idof_islandid_out: wp.array2d[int],
+  unconstrained_cnt_inout: wp.array2d[int],
+):
+  worldid, dofid = wp.tid()
+
+  nidof = nidof_in[worldid]
+  island_id = dof_island_in[worldid, dofid]
+
+  if island_id >= 0:
+    local_idx = wp.atomic_add(island_nv_inout, worldid, island_id, 1)
+    idof = island_idofadr_in[worldid, island_id] + local_idx
+    idof_islandid_out[worldid, idof] = island_id
+  else:
+    cnt = wp.atomic_add(unconstrained_cnt_inout, worldid, 0, 1)
+    idof = nidof + cnt
+
+  map_dof2idof_out[worldid, dofid] = idof
+  map_idof2dof_out[worldid, idof] = dofid
+
+
+@wp.kernel
+def _island_count_constraints(
+  nefc_in: wp.array[int],
+  njmax_in: int,
+  efc_tree_in: wp.array2d[int],
+  tree_island_in: wp.array2d[int],
+  efc_type_in: wp.array2d[int],
+  efc_island_out: wp.array2d[int],
+  island_nefc_out: wp.array2d[int],
+  island_ne_out: wp.array2d[int],
+  island_nf_out: wp.array2d[int],
+):
+  worldid, efcid = wp.tid()
+  if efcid >= wp.min(njmax_in, nefc_in[worldid]):
+    return
+
+  efc_tree = efc_tree_in[worldid, efcid]
+  if efc_tree < 0:
+    efc_island_out[worldid, efcid] = -1
+    return
+  island_id = tree_island_in[worldid, efc_tree]
+  efc_island_out[worldid, efcid] = island_id
+
+  if island_id >= 0:
+    wp.atomic_add(island_nefc_out, worldid, island_id, 1)
+
+    efc_type = efc_type_in[worldid, efcid]
+    if efc_type == ConstraintType.EQUALITY:
+      wp.atomic_add(island_ne_out, worldid, island_id, 1)
+    elif efc_type == ConstraintType.FRICTION_DOF or efc_type == ConstraintType.FRICTION_TENDON:
+      wp.atomic_add(island_nf_out, worldid, island_id, 1)
+
+
+@wp.kernel
+def _island_map_constraints(
+  nefc_in: wp.array[int],
+  njmax_in: int,
+  efc_island_in: wp.array2d[int],
+  island_iefcadr_in: wp.array2d[int],
+  island_ne_in: wp.array2d[int],
+  island_nf_in: wp.array2d[int],
+  efc_type_in: wp.array2d[int],
+  # Counters (inout):
+  island_ne_mapped_inout: wp.array2d[int],
+  island_nf_mapped_inout: wp.array2d[int],
+  island_nother_mapped_inout: wp.array2d[int],
+  island_nefc_inout: wp.array2d[int],
+  # Out:
+  map_efc2iefc_out: wp.array2d[int],
+  map_iefc2efc_out: wp.array2d[int],
+  iefc_islandid_out: wp.array2d[int],
+):
+  worldid, efcid = wp.tid()
+  if efcid >= wp.min(njmax_in, nefc_in[worldid]):
+    return
+
+  island_id = efc_island_in[worldid, efcid]
+  if island_id >= 0:
+    efc_type = efc_type_in[worldid, efcid]
+
+    # 1. Determine local index and absolute index ic based on category
+    if efc_type == ConstraintType.EQUALITY:
+      local_idx = wp.atomic_add(island_ne_mapped_inout, worldid, island_id, 1)
+      ic = island_iefcadr_in[worldid, island_id] + local_idx
+    elif efc_type == ConstraintType.FRICTION_DOF or efc_type == ConstraintType.FRICTION_TENDON:
+      local_idx = wp.atomic_add(island_nf_mapped_inout, worldid, island_id, 1)
+      ic = island_iefcadr_in[worldid, island_id] + island_ne_in[worldid, island_id] + local_idx
+    else:
+      local_idx = wp.atomic_add(island_nother_mapped_inout, worldid, island_id, 1)
+      ic = (
+        island_iefcadr_in[worldid, island_id] + island_ne_in[worldid, island_id] + island_nf_in[worldid, island_id] + local_idx
+      )
+
+    # 2. Increment overall island_nefc counter to reconstruct d.island_nefc
+    wp.atomic_add(island_nefc_inout, worldid, island_id, 1)
+
+    # 3. Store mappings
+    map_efc2iefc_out[worldid, efcid] = ic
+    map_iefc2efc_out[worldid, ic] = efcid
+    iefc_islandid_out[worldid, ic] = island_id
+
+
+@wp.kernel
+def _island_scan_sparse_rows(
+  nisland_in: wp.array[int],
+  efc_J_rownnz_in: wp.array2d[int],
+  island_nefc_in: wp.array2d[int],
+  island_iefcadr_in: wp.array2d[int],
+  map_iefc2efc_in: wp.array2d[int],
+  iefc_J_rownnz_out: wp.array2d[int],
+  iefc_J_rowadr_out: wp.array2d[int],
+):
+  worldid = wp.tid()
+
+  nisland = nisland_in[worldid]
+  if nisland == 0:
+    return
+
+  total_gathered_efc = island_iefcadr_in[worldid, nisland - 1] + island_nefc_in[worldid, nisland - 1]
+
+  running_rowadr = int(0)
+  for ic in range(total_gathered_efc):
+    c = map_iefc2efc_in[worldid, ic]
+    rownnz = efc_J_rownnz_in[worldid, c]
+    iefc_J_rowadr_out[worldid, ic] = running_rowadr
+    iefc_J_rownnz_out[worldid, ic] = rownnz
+    running_rowadr += rownnz
+
+
+@wp.kernel
+def _compute_efc_tree(
+  # Model:
+  nv: int,
+  body_treeid: wp.array[int],
+  jnt_dofadr: wp.array[int],
+  dof_treeid: wp.array[int],
+  geom_bodyid: wp.array[int],
+  site_bodyid: wp.array[int],
+  eq_type: wp.array[int],
+  eq_obj1id: wp.array[int],
+  eq_obj2id: wp.array[int],
+  eq_objtype: wp.array[int],
+  is_sparse: bool,
+  # Data in:
+  nefc_in: wp.array[int],
+  contact_geom_in: wp.array[wp.vec2i],
+  efc_type_in: wp.array2d[int],
+  efc_id_in: wp.array2d[int],
+  efc_J_in: wp.array3d[float],
+  efc_J_rownnz_in: wp.array2d[int],
+  efc_J_rowadr_in: wp.array2d[int],
+  efc_J_colind_in: wp.array3d[int],
+  njmax_in: int,
+  # Out:
+  efc_tree_out: wp.array2d[int],
+):
+  """Compute the first non-negative tree for each constraint."""
+  worldid, efcid = wp.tid()
+
+  if efcid >= wp.min(njmax_in, nefc_in[worldid]):
+    return
+
+  efc_type = efc_type_in[worldid, efcid]
+  efc_id = efc_id_in[worldid, efcid]
+
+  tree = int(-1)
+  use_generic = int(0)
+
+  # equality (connect/weld)
+  if efc_type == ConstraintType.EQUALITY:
+    eq_t = eq_type[efc_id]
+
+    if eq_t == EqType.CONNECT or eq_t == EqType.WELD:
+      b1 = eq_obj1id[efc_id]
+      b2 = eq_obj2id[efc_id]
+
+      if eq_objtype[efc_id] == ObjType.SITE:
+        b1 = site_bodyid[b1]
+        b2 = site_bodyid[b2]
+
+      t1 = body_treeid[b1]
+      t2 = body_treeid[b2]
+      if t1 >= 0:
+        tree = t1
+      else:
+        tree = t2
+    else:
+      # JOINT, TENDON, FLEX: generic scan
+      use_generic = 1
+
+  # joint friction
+  elif efc_type == ConstraintType.FRICTION_DOF:
+    tree = dof_treeid[efc_id]
+
+  # joint limit
+  elif efc_type == ConstraintType.LIMIT_JOINT:
+    tree = dof_treeid[jnt_dofadr[efc_id]]
+
+  # contact
+  elif (
+    efc_type == ConstraintType.CONTACT_FRICTIONLESS
+    or efc_type == ConstraintType.CONTACT_PYRAMIDAL
+    or efc_type == ConstraintType.CONTACT_ELLIPTIC
+  ):
+    geom_pair = contact_geom_in[efc_id]
+    g1 = geom_pair[0]
+    g2 = geom_pair[1]
+
+    if g1 >= 0 and g2 >= 0:
+      t1 = body_treeid[geom_bodyid[g1]]
+      t2 = body_treeid[geom_bodyid[g2]]
+      if t1 >= 0:
+        tree = t1
+      else:
+        tree = t2
+    else:
+      # flex contacts: generic scan
+      use_generic = 1
+
+  else:
+    # generic: scan Jacobian row
+    use_generic = 1
+
+  if use_generic:
+    count = nv
+    rowadr = 0
+    if is_sparse:
+      count = efc_J_rownnz_in[worldid, efcid]
+      rowadr = efc_J_rowadr_in[worldid, efcid]
+
+    for i in range(count):
+      dof = i
+      if is_sparse:
+        sparseid = rowadr + i
+        dof = efc_J_colind_in[worldid, 0, sparseid]
+      else:
+        J_val = efc_J_in[worldid, efcid, dof]
+        if J_val == 0.0:
+          continue
+
+      t = dof_treeid[dof]
+      if t >= 0:
+        tree = t
+        break
+
+  efc_tree_out[worldid, efcid] = tree
+
+
+@wp.kernel
+def _gather_efc_and_jacobian(
+  # Model:
+  is_sparse: bool,
+  # Data in:
+  nefc_in: wp.array[int],
+  efc_D_in: wp.array2d[float],
+  efc_type_in: wp.array2d[int],
+  efc_id_in: wp.array2d[int],
+  efc_frictionloss_in: wp.array2d[float],
+  efc_aref_in: wp.array2d[float],
+  efc_J_in: wp.array3d[float],
+  # Sparse Jacobian arrays:
+  efc_J_rownnz_in: wp.array2d[int],
+  efc_J_rowadr_in: wp.array2d[int],
+  efc_J_colind_in: wp.array3d[int],
+  iefc_J_rowadr_in: wp.array2d[int],
+  # In:
+  njmax_in: int,
+  map_iefc2efc_in: wp.array2d[int],
+  map_idof2dof_in: wp.array2d[int],
+  map_dof2idof_in: wp.array2d[int],
+  nidof_in: wp.array[int],
+  # Out:
+  iefc_D_out: wp.array2d[float],
+  iefc_type_out: wp.array2d[int],
+  iefc_id_out: wp.array2d[int],
+  iefc_frictionloss_out: wp.array2d[float],
+  iefc_aref_out: wp.array2d[float],
+  iefc_J_out: wp.array3d[float],
+  iefc_J_colind_out: wp.array3d[int],
+):
+  """Gather constraint arrays and Jacobian into island-local dense order."""
+  worldid, iefcid = wp.tid()
+
+  if iefcid >= wp.min(njmax_in, nefc_in[worldid]):
+    return
+
+  c = map_iefc2efc_in[worldid, iefcid]
+
+  # Scalar constraint fields
+  iefc_D_out[worldid, iefcid] = efc_D_in[worldid, c]
+  iefc_type_out[worldid, iefcid] = efc_type_in[worldid, c]
+  iefc_id_out[worldid, iefcid] = efc_id_in[worldid, c]
+  iefc_frictionloss_out[worldid, iefcid] = efc_frictionloss_in[worldid, c]
+  iefc_aref_out[worldid, iefcid] = efc_aref_in[worldid, c]
+
+  # Jacobian: convert to island-local dense format
+  nid = nidof_in[worldid]
+
+  if is_sparse:
+    rownnz = efc_J_rownnz_in[worldid, c]
+    rowadr_in = efc_J_rowadr_in[worldid, c]
+    rowadr_out = iefc_J_rowadr_in[worldid, iefcid]
+    for i in range(rownnz):
+      sparseid_in = rowadr_in + i
+      sparseid_out = rowadr_out + i
+      dof = efc_J_colind_in[worldid, 0, sparseid_in]
+      idof = map_dof2idof_in[worldid, dof]
+      # Store in sparse iefc_J_out and iefc_J_colind_out
+      iefc_J_out[worldid, 0, sparseid_out] = efc_J_in[worldid, 0, sparseid_in]
+      iefc_J_colind_out[worldid, 0, sparseid_out] = idof
+  else:
+    # Dense path: reorder rows by iefc, columns by idof
+    for idof in range(nid):
+      dof = map_idof2dof_in[worldid, idof]
+      iefc_J_out[worldid, iefcid, idof] = efc_J_in[worldid, c, dof]
+
+
+@wp.kernel
+def _gather_dof_arrays(
+  # Data in:
+  qacc_in: wp.array2d[float],
+  qacc_smooth_in: wp.array2d[float],
+  qfrc_smooth_in: wp.array2d[float],
+  # In:
+  nidof_in: wp.array[int],
+  map_idof2dof_in: wp.array2d[int],
+  # Out:
+  iacc_out: wp.array2d[float],
+  iacc_smooth_out: wp.array2d[float],
+  ifrc_smooth_out: wp.array2d[float],
+):
+  """Gather DOF arrays into island-local order."""
+  worldid, idofid = wp.tid()
+
+  if idofid >= nidof_in[worldid]:
+    return
+
+  dof = map_idof2dof_in[worldid, idofid]
+  iacc_out[worldid, idofid] = qacc_in[worldid, dof]
+  iacc_smooth_out[worldid, idofid] = qacc_smooth_in[worldid, dof]
+  ifrc_smooth_out[worldid, idofid] = qfrc_smooth_in[worldid, dof]
+
+
+@wp.kernel
+def _scatter_dof_arrays(
+  # In:
+  qacc_smooth_in: wp.array2d[float],
+  qfrc_smooth_in: wp.array2d[float],
+  dof_island_in: wp.array2d[int],
+  iacc_in: wp.array2d[float],
+  ifrc_constraint_in: wp.array2d[float],
+  iMa_in: wp.array2d[float],
+  map_dof2idof_in: wp.array2d[int],
+  scatter_Ma: bool,
+  # Out:
+  qacc_out: wp.array2d[float],
+  qfrc_constraint_out: wp.array2d[float],
+  Ma_out: wp.array2d[float],
+):
+  """Scatter island results to global arrays, copy qacc_smooth for non-island DOFs."""
+  worldid, dofid = wp.tid()
+
+  if dof_island_in[worldid, dofid] < 0:
+    qacc_out[worldid, dofid] = qacc_smooth_in[worldid, dofid]
+    qfrc_constraint_out[worldid, dofid] = 0.0
+    if scatter_Ma:
+      Ma_out[worldid, dofid] = qfrc_smooth_in[worldid, dofid]
+  else:
+    idof = map_dof2idof_in[worldid, dofid]
+    qacc_out[worldid, dofid] = iacc_in[worldid, idof]
+    qfrc_constraint_out[worldid, dofid] = ifrc_constraint_in[worldid, idof]
+    if scatter_Ma:
+      Ma_out[worldid, dofid] = iMa_in[worldid, idof]
+
+
+@wp.kernel
+def _scatter_efc_arrays(
+  # In:
+  nefc_in: wp.array[int],
+  njmax_in: int,
+  map_iefc2efc_in: wp.array2d[int],
+  iefc_force_in: wp.array2d[float],
+  iefc_state_in: wp.array2d[int],
+  # Out:
+  efc_force_out: wp.array2d[float],
+  efc_state_out: wp.array2d[int],
+):
+  """Scatter island-local constraint results back to global order."""
+  worldid, iefcid = wp.tid()
+
+  if iefcid >= wp.min(njmax_in, nefc_in[worldid]):
+    return
+
+  c = map_iefc2efc_in[worldid, iefcid]
+  efc_force_out[worldid, c] = iefc_force_in[worldid, iefcid]
+  efc_state_out[worldid, c] = iefc_state_in[worldid, iefcid]
+
+
+@wp.kernel
+def _init_island_arrays(
+  island_idofadr_out: wp.array2d[int],
+  island_nv_out: wp.array2d[int],
+  island_nefc_out: wp.array2d[int],
+  island_ne_out: wp.array2d[int],
+  island_nf_out: wp.array2d[int],
+  island_iefcadr_out: wp.array2d[int],
+  nidof_out: wp.array[int],
+):
+  worldid, islandid = wp.tid()
+  island_nv_out[worldid, islandid] = 0
+  island_nefc_out[worldid, islandid] = 0
+  island_ne_out[worldid, islandid] = 0
+  island_nf_out[worldid, islandid] = 0
+  island_idofadr_out[worldid, islandid] = 0
+  island_iefcadr_out[worldid, islandid] = 0
+  if islandid == 0:
+    nidof_out[worldid] = 0
+
+
+@wp.kernel
+def _init_dof_arrays(
+  dof_island_out: wp.array2d[int],
+  map_dof2idof_out: wp.array2d[int],
+  map_idof2dof_out: wp.array2d[int],
+  idof_islandid_out: wp.array2d[int],
+):
+  worldid, dofid = wp.tid()
+  dof_island_out[worldid, dofid] = -1
+  map_dof2idof_out[worldid, dofid] = 0
+  map_idof2dof_out[worldid, dofid] = 0
+  idof_islandid_out[worldid, dofid] = -1
+
+
+@wp.kernel
+def _init_efc_arrays(
+  efc_island_out: wp.array2d[int],
+  map_efc2iefc_out: wp.array2d[int],
+  map_iefc2efc_out: wp.array2d[int],
+  iefc_islandid_out: wp.array2d[int],
+  efc_tree_out: wp.array2d[int],
+):
+  worldid, efcid = wp.tid()
+  efc_island_out[worldid, efcid] = -1
+  map_efc2iefc_out[worldid, efcid] = 0
+  map_iefc2efc_out[worldid, efcid] = 0
+  iefc_islandid_out[worldid, efcid] = -1
+  efc_tree_out[worldid, efcid] = -1
+
+
+@event_scope
+def compute_island_mapping(m: types.Model, d: types.Data, ctx: IslandSolverContext):
+  """Compute DOF/constraint island mappings after island discovery.
+
+  Populates island solver context arrays via ctx: nv, nefc, ne, nf,
+  iefcadr, nidof, map_dof2idof, map_idof2dof, dof_islandid, map_efc2iefc,
+  map_iefc2efc, efc_islandid. Also populates d.dof_island, d.efc.island,
+  and d.island_dofadr.
+
+  Args:
+    m: Model.
+    d: Data.
+    ctx: IslandSolverContext.
+  """
+  # Ensure dof_islandid / efc_islandid are allocated at the right shape
+  if d.dof_islandid.shape[1] != m.nv:
+    d.dof_islandid = wp.empty((d.nworld, m.nv), dtype=int)
+  if d.efc_islandid.shape[1] != d.njmax:
+    d.efc_islandid = wp.empty((d.nworld, d.njmax), dtype=int)
+
+  # Ensure island-local DOF arrays are allocated at the right shape
+  if d.iqacc.shape[1] != m.nv:
+    nw = d.nworld
+    d.iqacc = wp.empty((nw, m.nv), dtype=float)
+    d.iqacc_smooth = wp.empty((nw, m.nv), dtype=float)
+    d.iqfrc_smooth = wp.empty((nw, m.nv), dtype=float)
+    d.iqfrc_constraint = wp.empty((nw, m.nv), dtype=float)
+  wp.launch(
+    _init_island_arrays,
+    dim=(d.nworld, m.ntree),
+    inputs=[],
+    outputs=[
+      d.island_dofadr,
+      d.island_nv,
+      d.island_nefc,
+      d.island_ne,
+      d.island_nf,
+      d.island_efcadr,
+      d.nidof,
+    ],
+  )
+  wp.launch(
+    _init_dof_arrays,
+    dim=(d.nworld, m.nv),
+    inputs=[],
+    outputs=[d.dof_island, d.map_dof2idof, d.map_idof2dof, d.dof_islandid],
+  )
+  efc_tree = wp.empty((d.nworld, d.njmax), dtype=int)
+  wp.launch(
+    _init_efc_arrays,
+    dim=(d.nworld, d.njmax),
+    inputs=[],
+    outputs=[d.efc.island, d.map_efc2iefc, d.map_iefc2efc, d.efc_islandid, efc_tree],
+  )
 
   wp.launch(
-    _flood_fill,
-    dim=d.nworld,
-    inputs=[m.ntree, tree_tree, d.tree_island, stack_scratch],
-    outputs=[d.nisland, d.tree_island, stack_scratch],
+    _compute_efc_tree,
+    dim=(d.nworld, d.njmax),
+    inputs=[
+      m.nv,
+      m.body_treeid,
+      m.jnt_dofadr,
+      m.dof_treeid,
+      m.geom_bodyid,
+      m.site_bodyid,
+      m.eq_type,
+      m.eq_obj1id,
+      m.eq_obj2id,
+      m.eq_objtype,
+      m.is_sparse,
+      d.nefc,
+      d.contact.geom,
+      d.efc.type,
+      d.efc.id,
+      d.efc.J,
+      d.efc.J_rownnz,
+      d.efc.J_rowadr,
+      d.efc.J_colind,
+      d.njmax,
+    ],
+    outputs=[efc_tree],
+  )
+
+  # 1. Count DOFs per island
+  wp.launch(
+    _island_count_dofs,
+    dim=(d.nworld, m.nv),
+    inputs=[m.dof_treeid, d.tree_island],
+    outputs=[d.dof_island, d.island_nv],
+  )
+
+  # 2. Count Constraints per island
+  wp.launch(
+    _island_count_constraints,
+    dim=(d.nworld, d.njmax),
+    inputs=[d.nefc, d.njmax, efc_tree, d.tree_island, d.efc.type],
+    outputs=[d.efc.island, d.island_nefc, d.island_ne, d.island_nf],
+  )
+
+  # 3. Scan sizes and reset counters for mapping
+  wp.launch(
+    _island_scan_sizes,
+    dim=d.nworld,
+    inputs=[d.nisland],
+    outputs=[d.island_dofadr, d.island_nv, d.island_nefc, d.island_efcadr, d.nidof],
+  )
+
+  # 4. Map DOFs
+  unconstrained_cnt = wp.zeros((d.nworld, 1), dtype=int)
+  wp.launch(
+    _island_map_dofs,
+    dim=(d.nworld, m.nv),
+    inputs=[m.nv, d.dof_island, d.island_dofadr, d.nidof],
+    outputs=[d.island_nv, d.map_dof2idof, d.map_idof2dof, d.dof_islandid, unconstrained_cnt],
+  )
+
+  # 5. Map Constraints
+  ne_mapped = wp.zeros((d.nworld, m.ntree), dtype=int)
+  nf_mapped = wp.zeros((d.nworld, m.ntree), dtype=int)
+  nother_mapped = wp.zeros((d.nworld, m.ntree), dtype=int)
+
+  wp.launch(
+    _island_map_constraints,
+    dim=(d.nworld, d.njmax),
+    inputs=[
+      d.nefc,
+      d.njmax,
+      d.efc.island,
+      d.island_efcadr,
+      d.island_ne,
+      d.island_nf,
+      d.efc.type,
+      ne_mapped,
+      nf_mapped,
+      nother_mapped,
+    ],
+    outputs=[d.island_nefc, d.map_efc2iefc, d.map_iefc2efc, d.efc_islandid],
+  )
+
+  # 6. Scan Sparse Rows (if sparse)
+  if m.is_sparse:
+    wp.launch(
+      _island_scan_sparse_rows,
+      dim=d.nworld,
+      inputs=[d.nisland, d.efc.J_rownnz, d.island_nefc, d.island_efcadr, d.map_iefc2efc],
+      outputs=[d.efc.iJ_rownnz, d.efc.iJ_rowadr],
+    )
+
+
+@event_scope
+def gather_island_inputs(m: types.Model, d: types.Data, ctx: IslandSolverContext):
+  """Gather constraint and DOF arrays into island-local order.
+
+  Populates d.iefc (D, type, id, frictionloss, aref, J, J_colind) and
+  d.iqacc, d.iqacc_smooth, d.iqfrc_smooth.
+
+  Must be called after compute_island_mapping() and before per-island solving.
+
+  Args:
+    m: Model.
+    d: Data.
+    ctx: IslandSolverContext whose arrays are populated.
+  """
+  # Gather constraint arrays and dense Jacobian (fused)
+  wp.launch(
+    _gather_efc_and_jacobian,
+    dim=(d.nworld, d.njmax),
+    inputs=[
+      m.is_sparse,
+      d.nefc,
+      d.efc.D,
+      d.efc.type,
+      d.efc.id,
+      d.efc.frictionloss,
+      d.efc.aref,
+      d.efc.J,
+      d.efc.J_rownnz,
+      d.efc.J_rowadr,
+      d.efc.J_colind,
+      d.efc.iJ_rowadr,
+      d.njmax,
+      d.map_iefc2efc,
+      d.map_idof2dof,
+      d.map_dof2idof,
+      d.nidof,
+    ],
+    outputs=[
+      d.efc.iD,
+      d.efc.itype,
+      d.efc.iid,
+      d.efc.ifrictionloss,
+      d.efc.iaref,
+      d.efc.iJ,
+      d.efc.iJ_colind,
+    ],
+  )
+
+  # Gather DOF arrays
+  wp.launch(
+    _gather_dof_arrays,
+    dim=(d.nworld, m.nv),
+    inputs=[
+      d.qacc,
+      d.qacc_smooth,
+      d.qfrc_smooth,
+      d.nidof,
+      d.map_idof2dof,
+    ],
+    outputs=[
+      d.iqacc,
+      d.iqacc_smooth,
+      d.iqfrc_smooth,
+    ],
+  )
+
+
+@event_scope
+def scatter_island_results(m: types.Model, d: types.Data, ctx: IslandSolverContext, scatter_Ma: bool):
+  """Scatter island-local solver results back to global arrays.
+
+  Reads ctx qacc, qfrc_constraint, Ma, d.efc iforce, istate and
+  writes them back to d.qacc, d.qfrc_constraint, d.efc.Ma, d.efc.force,
+  d.efc.state. Unconstrained DOFs receive qacc_smooth and zero qfrc.
+
+  Args:
+    m: Model.
+    d: Data.
+    ctx: IslandSolverContext which contains results.
+    scatter_Ma: Whether to scatter Ma for Euler/implicit integrators.
+  """
+  # Scatter DOF results (and optionally Ma for Euler/implicit integrators)
+  wp.launch(
+    _scatter_dof_arrays,
+    dim=(d.nworld, m.nv),
+    inputs=[
+      d.qacc_smooth,
+      d.qfrc_smooth,
+      d.dof_island,
+      d.iqacc,
+      d.iqfrc_constraint,
+      ctx.Ma,
+      d.map_dof2idof,
+      scatter_Ma,
+    ],
+    outputs=[
+      d.qacc,
+      d.qfrc_constraint,
+      d.efc.Ma,
+    ],
+  )
+
+  # Scatter constraint results (force and state from island-local arrays)
+  wp.launch(
+    _scatter_efc_arrays,
+    dim=(d.nworld, d.njmax),
+    inputs=[
+      d.nefc,
+      d.njmax,
+      d.map_iefc2efc,
+      d.efc.iforce,
+      d.efc.istate,
+    ],
+    outputs=[
+      d.efc.force,
+      d.efc.state,
+    ],
   )
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/ray.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/ray.py
index a320644a..d9b20861 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/ray.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/ray.py
@@ -707,10 +707,14 @@ def ray_mesh_with_bvh(
   pnt: wp.vec3,
   vec: wp.vec3,
   max_t: float,
+  cull_backfaces: bool,
 ) -> Tuple[float, wp.vec3, float, float, int, int]:
   """Returns intersection information for ray mesh intersections.
 
   Requires wp.Mesh be constructed and their ids to be passed.
+
+  When ``cull_backfaces`` is True, the function rejects exit-face hits in the
+  local-space frame. This matches MuJoCo OpenGL rendering's backface culling rule.
   """
   t = float(-1.0)
   u = float(0.0)
@@ -722,7 +726,12 @@ def ray_mesh_with_bvh(
   lpnt, lvec = _ray_map(pos, mat, pnt, vec)
   hit = wp.mesh_query_ray(mesh_bvh_id[mesh_geom_id], lpnt, lvec, max_t, t, u, v, sign, n, f)
 
-  if hit and wp.dot(lvec, n) < 0.0:  # Backface culling in local space
+  if not hit:
+    return -1.0, wp.vec3(0.0, 0.0, 0.0), 0.0, 0.0, -1, -1
+
+  # Front-face hit, or back-face hit when cull is disabled: rotate the
+  # local-space normal into world space and return the hit.
+  if (not cull_backfaces) or wp.dot(lvec, n) < 0.0:
     normal = mat @ n
     normal = wp.normalize(normal)
     return t, normal, u, v, f, mesh_geom_id
@@ -1055,6 +1064,8 @@ def _ray_geom_mesh_bvh(
 
     if gtype == GeomType.MESH or gtype == GeomType.HFIELD:
       bvh_ids = mesh_bvh_id if gtype == GeomType.MESH else hfield_bvh_id
+      # Public ray API (mjw.ray / mjw.rays) preserves MuJoCo's mj_ray cull rule:
+      # rangefinder sensors and user-facing ray casts always cull back-faces.
       t, n, u, v, f, geom_mesh_id = ray_mesh_with_bvh(
         bvh_ids,
         geom_dataid[worldid % geom_dataid.shape[0], geomid],
@@ -1063,6 +1074,7 @@ def _ray_geom_mesh_bvh(
         pnt,
         vec,
         min_dist,
+        True,
       )
       if t >= 0.0 and t < min_dist:
         return t, n
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/render.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/render.py
index ba3b4cfe..7afb140b 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/render.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/render.py
@@ -180,6 +180,7 @@ def cast_ray(
   flex_group_root: wp.array2d[int],
   ray_origin_world: wp.vec3,
   ray_dir_world: wp.vec3,
+  cull_backfaces: bool,
 ) -> Tuple[int, float, wp.vec3, float, float, int, int]:
   dist = float(MJ_MAXVAL)
   normal = wp.vec3(0.0, 0.0, 0.0)
@@ -232,6 +233,7 @@ def cast_ray(
         ray_origin_world,
         ray_dir_world,
         dist,
+        cull_backfaces,
       )
     if gtype == GeomType.SPHERE:
       d, n = ray_sphere(
@@ -281,6 +283,7 @@ def cast_ray(
         ray_origin_world,
         ray_dir_world,
         dist,
+        cull_backfaces,
       )
     if gtype == GeomType.FLEX:
       hit_geom_id = -2
@@ -308,6 +311,11 @@ def cast_ray(
         if d >= 0.0:
           hit_mesh_id = flexid
 
+    # Backface cull: drop exit-face hits when the ray origin is inside the geom,
+    # matching ray_mesh_with_bvh's `dot(lvec, n) < 0` rule.
+    if cull_backfaces and d >= 0.0 and wp.dot(ray_dir_world, n) > 0.0:
+      d = -1.0
+
     if d >= 0.0 and d < dist:
       dist = d
       normal = n
@@ -349,6 +357,7 @@ def cast_ray_first_hit(
   ray_origin_world: wp.vec3,
   ray_dir_world: wp.vec3,
   max_dist: float,
+  cull_backfaces: bool,
 ) -> bool:
   """A simpler version of casting rays that only checks for the first hit."""
   query = wp.bvh_query_ray(bvh_id, ray_origin_world, ray_dir_world, group_root)
@@ -387,6 +396,7 @@ def cast_ray_first_hit(
         ray_origin_world,
         ray_dir_world,
         max_dist,
+        cull_backfaces,
       )
     if gtype == GeomType.SPHERE:
       d, n = ray_sphere(
@@ -467,6 +477,11 @@ def cast_ray_first_hit(
         )
         d = 0.0 if hit else -1.0
 
+    # Backface cull: see cast_ray for rationale. Strict `> 0` keeps tangent
+    # hits and skips branches with a zero-vector normal (mesh/flex anyhit).
+    if cull_backfaces and d >= 0.0 and wp.dot(ray_dir_world, n) > 0.0:
+      d = -1.0
+
     if d >= 0.0 and d < max_dist:
       return True
 
@@ -507,6 +522,7 @@ def compute_lighting(
   lightdir: wp.vec3,
   normal: wp.vec3,
   hitpoint: wp.vec3,
+  cull_backfaces: bool,
 ) -> float:
   light_contribution = float(0.0)
 
@@ -571,6 +587,7 @@ def compute_lighting(
       shadow_origin,
       L,
       max_t,
+      cull_backfaces,
     )
 
     if shadow_hit:
@@ -725,6 +742,7 @@ def render(m: Model, d: Data, rc: RenderContext):
       flex_group_root,
       ray_origin_world,
       ray_dir_world,
+      wp.static(rc.enable_backface_culling),
     )
 
     if render_seg[cam_idx] and geom_id != -1:
@@ -798,12 +816,14 @@ def render(m: Model, d: Data, rc: RenderContext):
             )
             base_color = wp.cw_mul(base_color, tex_color)
 
-    len_n = wp.length(normal)
-    n = normal if len_n > 0.0 else wp.vec3(0.0, 0.0, 1.0)
-    n = wp.normalize(n)
-    hemispheric = 0.5 * (n[2] + 1.0)
-    ambient_color = wp.vec3(0.4, 0.4, 0.45) * hemispheric + wp.vec3(0.1, 0.1, 0.12) * (1.0 - hemispheric)
-    result = 0.5 * wp.cw_mul(base_color, ambient_color)
+    result = wp.vec3(0.0, 0.0, 0.0)
+    if wp.static(rc.use_ambient_lighting):
+      len_n = wp.length(normal)
+      n = normal if len_n > 0.0 else wp.vec3(0.0, 0.0, 1.0)
+      n = wp.normalize(n)
+      hemispheric = 0.5 * (n[2] + 1.0)
+      ambient_color = wp.vec3(0.4, 0.4, 0.45) * hemispheric + wp.vec3(0.1, 0.1, 0.12) * (1.0 - hemispheric)
+      result = 0.5 * wp.cw_mul(base_color, ambient_color)
 
     # Apply lighting and shadows
     for l in range(wp.static(m.nlight)):
@@ -837,6 +857,7 @@ def render(m: Model, d: Data, rc: RenderContext):
         light_xdir_in[worldid, l],
         normal,
         hit_point,
+        wp.static(rc.enable_backface_culling),
       )
       result = result + base_color * light_contribution
 
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py
index 472fa617..619e55cf 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py
@@ -17,6 +17,7 @@ from typing import Any, Tuple
 
 import warp as wp
 
+from mujoco.mjx.third_party.mujoco_warp._src import history
 from mujoco.mjx.third_party.mujoco_warp._src import math
 from mujoco.mjx.third_party.mujoco_warp._src import ray
 from mujoco.mjx.third_party.mujoco_warp._src import smooth
@@ -26,6 +27,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.collision_sdf import sdf
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXCONPAIR
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
+from mujoco.mjx.third_party.mujoco_warp._src.types import TACTILE_DEPTH_SEMANTICS
 from mujoco.mjx.third_party.mujoco_warp._src.types import ConeType
 from mujoco.mjx.third_party.mujoco_warp._src.types import ConstraintType
 from mujoco.mjx.third_party.mujoco_warp._src.types import ContactType
@@ -917,6 +919,10 @@ def sensor_pos(m: Model, d: Data):
     ],
   )
 
+  # apply sensor delay/interval for position sensors
+  history.apply_sensor_delay(m, d, m.sensor_pos_adr)
+  history.apply_sensor_delay(m, d, m.sensor_limitpos_adr)
+
   if m.callback.sensor:
     m.callback.sensor(m, d, Stage.POS)
 
@@ -1459,6 +1465,10 @@ def sensor_vel(m: Model, d: Data):
     ],
   )
 
+  # apply sensor delay/interval for velocity sensors
+  history.apply_sensor_delay(m, d, m.sensor_vel_adr)
+  history.apply_sensor_delay(m, d, m.sensor_limitvel_adr)
+
   if m.callback.sensor:
     m.callback.sensor(m, d, Stage.VEL)
 
@@ -2253,13 +2263,23 @@ def _sensor_tactile(
     vel_rel = vel_sensor - vel_other
 
     forceT = wp.vec3(0.0, 0.0, 0.0)
-    forceT[0] = -depth
+    if wp.static(TACTILE_DEPTH_SEMANTICS):
+      forceT[0] = -depth
+    else:
+      kMaxDepth = 0.05
+      pressure = depth / wp.max(kMaxDepth - depth, MJ_MINVAL)
+      force = wp.mul(normal, pressure)
+      forceT[0] = wp.dot(force, normal)
+
     if has_frame:
       forceT[1] = wp.abs(wp.dot(vel_rel, tang1))
       forceT[2] = wp.abs(wp.dot(vel_rel, tang2))
 
     dim = sensor_dim[sensor_id] // 3
-    wp.atomic_max(sensordata_out, worldid, sensor_adr[sensor_id] + 0 * dim + vertid, forceT[0])
+    if wp.static(TACTILE_DEPTH_SEMANTICS):
+      wp.atomic_max(sensordata_out, worldid, sensor_adr[sensor_id] + 0 * dim + vertid, forceT[0])
+    else:
+      wp.atomic_add(sensordata_out, worldid, sensor_adr[sensor_id] + 0 * dim + vertid, forceT[0])
     wp.atomic_add(sensordata_out, worldid, sensor_adr[sensor_id] + 1 * dim + vertid, forceT[1])
     wp.atomic_add(sensordata_out, worldid, sensor_adr[sensor_id] + 2 * dim + vertid, forceT[2])
 
@@ -2705,6 +2725,10 @@ def sensor_acc(m: Model, d: Data):
     ],
   )
 
+  # apply sensor delay/interval for acceleration sensors
+  history.apply_sensor_delay(m, d, m.sensor_acc_adr)
+  history.apply_sensor_delay(m, d, m.sensor_limitfrc_adr)
+
   if m.callback.sensor:
     m.callback.sensor(m, d, Stage.ACC)
 
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/smooth.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/smooth.py
index bda36c7c..4f35eab0 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/smooth.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/smooth.py
@@ -823,37 +823,38 @@ def _crb_accumulate(
 
 
 @wp.kernel
-def _qM_sparse(
+def _M_sparse(
   # Model:
   dof_bodyid: wp.array[int],
   dof_parentid: wp.array[int],
-  dof_Madr: wp.array[int],
   dof_armature: wp.array2d[float],
+  M_rownnz: wp.array[int],
+  M_rowadr: wp.array[int],
   # Data in:
   cdof_in: wp.array2d[wp.spatial_vector],
   crb_in: wp.array2d[vec10],
   # Data out:
-  qM_out: wp.array3d[float],
+  M_out: wp.array3d[float],
 ):
   worldid, dofid = wp.tid()
-  madr_ij = dof_Madr[dofid]  # dof_Madr is not batched
   bodyid = dof_bodyid[dofid]
+  madr_ij = M_rowadr[dofid] + M_rownnz[dofid] - 1
 
   # init M(i,i) with armature inertia
-  qM_out[worldid, 0, madr_ij] = dof_armature[worldid % dof_armature.shape[0], dofid]
+  M_out[worldid, 0, madr_ij] = dof_armature[worldid % dof_armature.shape[0], dofid]
 
   # precompute buf = crb_body_i * cdof_i
   buf = math.inert_vec(crb_in[worldid, bodyid], cdof_in[worldid, dofid])
 
   # sparse backward pass over ancestors
   while dofid >= 0:
-    qM_out[worldid, 0, madr_ij] += wp.dot(cdof_in[worldid, dofid], buf)
-    madr_ij += 1
+    M_out[worldid, 0, madr_ij] += wp.dot(cdof_in[worldid, dofid], buf)
+    madr_ij -= 1
     dofid = dof_parentid[dofid]
 
 
 @wp.kernel
-def _qM_dense(
+def _M_dense(
   # Model:
   dof_bodyid: wp.array[int],
   dof_parentid: wp.array[int],
@@ -862,7 +863,7 @@ def _qM_dense(
   cdof_in: wp.array2d[wp.spatial_vector],
   crb_in: wp.array2d[vec10],
   # Data out:
-  qM_out: wp.array3d[float],
+  M_out: wp.array3d[float],
 ):
   worldid, dofid = wp.tid()
   bodyid = dof_bodyid[dofid]
@@ -873,15 +874,15 @@ def _qM_dense(
   buf = math.inert_vec(crb_in[worldid, bodyid], cdof_in[worldid, dofid])
   M += wp.dot(cdof_in[worldid, dofid], buf)
 
-  qM_out[worldid, dofid, dofid] = M
+  M_out[worldid, dofid, dofid] = M
 
   # sparse backward pass over ancestors
   dofidi = dofid
   dofid = dof_parentid[dofid]
   while dofid >= 0:
-    qMij = wp.dot(cdof_in[worldid, dofid], buf)
-    qM_out[worldid, dofidi, dofid] += qMij
-    qM_out[worldid, dofid, dofidi] += qMij
+    Mij = wp.dot(cdof_in[worldid, dofid], buf)
+    M_out[worldid, dofidi, dofid] += Mij
+    M_out[worldid, dofid, dofidi] += Mij
     dofid = dof_parentid[dofid]
 
 
@@ -898,17 +899,17 @@ def crb(m: Model, d: Data):
     body_tree = m.body_tree[i]
     wp.launch(_crb_accumulate, dim=(d.nworld, body_tree.size), inputs=[m.body_parentid, d.crb, body_tree], outputs=[d.crb])
 
-  d.qM.zero_()
+  d.M.zero_()
   if m.is_sparse:
     wp.launch(
-      _qM_sparse,
+      _M_sparse,
       dim=(d.nworld, m.nv),
-      inputs=[m.dof_bodyid, m.dof_parentid, m.dof_Madr, m.dof_armature, d.cdof, d.crb],
-      outputs=[d.qM],
+      inputs=[m.dof_bodyid, m.dof_parentid, m.dof_armature, m.M_rownnz, m.M_rowadr, d.cdof, d.crb],
+      outputs=[d.M],
     )
   else:
     wp.launch(
-      _qM_dense, dim=(d.nworld, m.nv), inputs=[m.dof_bodyid, m.dof_parentid, m.dof_armature, d.cdof, d.crb], outputs=[d.qM]
+      _M_dense, dim=(d.nworld, m.nv), inputs=[m.dof_bodyid, m.dof_parentid, m.dof_armature, d.cdof, d.crb], outputs=[d.M]
     )
 
 
@@ -916,16 +917,17 @@ def crb(m: Model, d: Data):
 def _tendon_armature(
   # Model:
   dof_parentid: wp.array[int],
-  dof_Madr: wp.array[int],
   ten_J_rownnz: wp.array[int],
   ten_J_rowadr: wp.array[int],
   ten_J_colind: wp.array[int],
   tendon_armature: wp.array2d[float],
+  M_rownnz: wp.array[int],
+  M_rowadr: wp.array[int],
   is_sparse: bool,
   # Data in:
   ten_J_in: wp.array2d[float],
   # Data out:
-  qM_out: wp.array3d[float],
+  M_out: wp.array3d[float],
 ):
   worldid, tenid, dofid = wp.tid()
 
@@ -947,7 +949,7 @@ def _tendon_armature(
     return
 
   if is_sparse:
-    madr_ij = dof_Madr[dofid]
+    madr_ij = M_rowadr[dofid] + M_rownnz[dofid] - 1
 
   # sparse backward pass over ancestors
   dofidi = dofid
@@ -967,52 +969,40 @@ def _tendon_armature(
       else:
         ten_Jj = float(0.0)
 
-    qMij = armature * ten_Jj * ten_Ji
+    Mij = armature * ten_Jj * ten_Ji
 
     if is_sparse:
-      wp.atomic_add(qM_out[worldid, 0], madr_ij, qMij)
-      madr_ij += 1
+      wp.atomic_add(M_out[worldid, 0], madr_ij, Mij)
+      madr_ij -= 1
     else:
-      wp.atomic_add(qM_out[worldid, dofidi], dofid, qMij)
+      wp.atomic_add(M_out[worldid, dofidi], dofid, Mij)
       if dofidi != dofid:
-        wp.atomic_add(qM_out[worldid, dofid], dofidi, qMij)
+        wp.atomic_add(M_out[worldid, dofid], dofidi, Mij)
 
     dofid = dof_parentid[dofid]
 
 
 @event_scope
 def tendon_armature(m: Model, d: Data):
-  """Add tendon armature to qM."""
+  """Add tendon armature to M."""
   wp.launch(
     _tendon_armature,
     dim=(d.nworld, m.ntendon, m.max_ten_J_rownnz),
     inputs=[
       m.dof_parentid,
-      m.dof_Madr,
       m.ten_J_rownnz,
       m.ten_J_rowadr,
       m.ten_J_colind,
       m.tendon_armature,
+      m.M_rownnz,
+      m.M_rowadr,
       m.is_sparse,
       d.ten_J,
     ],
-    outputs=[d.qM],
+    outputs=[d.M],
   )
 
 
-@wp.kernel
-def _copy_CSR(
-  # Model:
-  mapM2M: wp.array[int],
-  # In:
-  M_in: wp.array3d[float],
-  # Out:
-  L_out: wp.array3d[float],
-):
-  worldid, ind = wp.tid()
-  L_out[worldid, 0, ind] = M_in[worldid, 0, mapM2M[ind]]
-
-
 @wp.kernel
 def _qLD_acc(
   # Model:
@@ -1055,7 +1045,7 @@ def _qLDiag_div(
 
 def _factor_i_sparse(m: Model, d: Data, M: wp.array3d[float], L: wp.array3d[float], D: wp.array2d[float]):
   """Sparse L'*D*L factorization of inertia-like matrix M, assumed spd."""
-  wp.launch(_copy_CSR, dim=(d.nworld, m.nC), inputs=[m.mapM2M, M], outputs=[L])
+  wp.copy(L, M)
 
   for i in reversed(range(len(m.qLD_updates))):
     qLD_updates = m.qLD_updates[i]
@@ -1071,7 +1061,7 @@ def _tile_cholesky_factorize(tile: TileSet):
   @wp.kernel(module="unique", enable_backward=False)
   def cholesky_factorize(
     # Data in:
-    qM_in: wp.array3d[float],
+    M_in: wp.array3d[float],
     # In:
     adr: wp.array[int],
     # Out:
@@ -1081,8 +1071,8 @@ def _tile_cholesky_factorize(tile: TileSet):
     TILE_SIZE = wp.static(tile.size)
 
     dofid = adr[nodeid]
-    M_tile = wp.tile_load(qM_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid))
-    L_tile = wp.tile_cholesky(M_tile)
+    M_tile = wp.tile_load(M_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid))
+    L_tile = wp.tile_cholesky(M_tile, fill_mode="upper")
     wp.tile_store(L_out[worldid], L_tile, offset=(dofid, dofid))
 
   return cholesky_factorize
@@ -1090,7 +1080,7 @@ def _tile_cholesky_factorize(tile: TileSet):
 
 def _factor_i_dense(m: Model, d: Data, M: wp.array, L: wp.array):
   """Dense Cholesky factorization of inertia-like matrix M, assumed spd."""
-  for tile in m.qM_tiles:
+  for tile in m.M_tiles:
     wp.launch_tiled(
       _tile_cholesky_factorize(tile),
       dim=(d.nworld, tile.adr.size),
@@ -1104,9 +1094,9 @@ def _factor_i_dense(m: Model, d: Data, M: wp.array, L: wp.array):
 def factor_m(m: Model, d: Data):
   """Factorization of inertia-like matrix M, assumed spd."""
   if m.is_sparse:
-    _factor_i_sparse(m, d, d.qM, d.qLD, d.qLDiagInv)
+    _factor_i_sparse(m, d, d.M, d.qLD, d.qLDiagInv)
   else:
-    _factor_i_dense(m, d, d.qM, d.qLD)
+    _factor_i_dense(m, d, d.M, d.qLD)
 
 
 @wp.kernel
@@ -2842,15 +2832,15 @@ def _tile_cholesky_solve(tile: TileSet):
     dofid = adr[nodeid]
     y_slice = wp.tile_load(y[worldid], shape=(TILE_SIZE,), offset=(dofid,))
     L_tile = wp.tile_load(L[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid))
-    x_slice = wp.tile_cholesky_solve(L_tile, y_slice)
+    x_slice = wp.tile_cholesky_solve(L_tile, y_slice, fill_mode="upper")
     wp.tile_store(x[worldid], x_slice, offset=(dofid,))
 
   return cholesky_solve
 
 
 def _solve_LD_dense(m: Model, d: Data, L: wp.array3d[float], x: wp.array2d[float], y: wp.array2d[float]):
-  """Computes dense backsubstitution: x = inv(L'*L)*y."""
-  for tile in m.qM_tiles:
+  """Computes dense backsubstitution: x = inv(U.T @ U) * y."""
+  for tile in m.M_tiles:
     wp.launch_tiled(
       _tile_cholesky_solve(tile),
       dim=(d.nworld, tile.adr.size),
@@ -2868,9 +2858,9 @@ def solve_LD(
   x: wp.array2d[float],
   y: wp.array2d[float],
 ):
-  """Computes backsubstitution to solve a linear system of the form x = inv(L'*D*L) * y.
+  """Computes backsubstitution for the inertia factorization.
 
-  L and D are the factors from the Cholesky factorization of the inertia matrix.
+  Sparse models use MuJoCo's L'*D*L factors; dense models use an upper Cholesky factor U.
 
   This function dispatches to either a sparse or dense solver depending on Model options.
 
@@ -2907,8 +2897,9 @@ def _tile_cholesky_factorize_solve(tile: TileSet):
 
   @wp.kernel(module="unique", enable_backward=False)
   def cholesky_factorize_solve(
+    # Data in:
+    M_in: wp.array3d[float],
     # In:
-    M: wp.array3d[float],
     y: wp.array2d[float],
     adr: wp.array[int],
     # Out:
@@ -2919,12 +2910,12 @@ def _tile_cholesky_factorize_solve(tile: TileSet):
     TILE_SIZE = wp.static(tile.size)
 
     dofid = adr[nodeid]
-    M_tile = wp.tile_load(M[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid))
+    M_tile = wp.tile_load(M_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid))
     y_slice = wp.tile_load(y[worldid], shape=(TILE_SIZE,), offset=(dofid,))
 
-    L_tile = wp.tile_cholesky(M_tile)
+    L_tile = wp.tile_cholesky(M_tile, fill_mode="upper")
     wp.tile_store(L[worldid], L_tile, offset=(dofid, dofid))
-    x_slice = wp.tile_cholesky_solve(L_tile, y_slice)
+    x_slice = wp.tile_cholesky_solve(L_tile, y_slice, fill_mode="upper")
     wp.tile_store(x[worldid], x_slice, offset=(dofid,))
 
   return cholesky_factorize_solve
@@ -2938,7 +2929,7 @@ def _factor_solve_i_dense(
   y: wp.array2d[float],
   L: wp.array3d[float],
 ):
-  for tile in m.qM_tiles:
+  for tile in m.M_tiles:
     wp.launch_tiled(
       _tile_cholesky_factorize_solve(tile),
       dim=(d.nworld, tile.adr.size),
@@ -2949,9 +2940,9 @@ def _factor_solve_i_dense(
 
 
 def factor_solve_i(m, d, M, L, D, x, y):
-  """Factorizes and solves the linear system: x = inv(L'*D*L) * y or x = inv(L'*L) * y.
+  """Factorizes and solves the inertia-like linear system.
 
-  M is an inertia-like matrix and L, D are its Cholesky-like factors.
+  Sparse models use MuJoCo's L'*D*L factors; dense models use an upper Cholesky factor U.
 
   This function first factorizes the matrix M (sparse or dense depending on model options),
   then solves the system for x given right-hand side y.
@@ -2960,7 +2951,7 @@ def factor_solve_i(m, d, M, L, D, x, y):
     m: The model containing factorization and sparsity information.
     d: The data object containing workspace and factorization results.
     M: The inertia-like matrix to factorize.
-    L: Output lower-triangular factor from the factorization (sparse or dense).
+    L: Output sparse factor or dense upper Cholesky factor.
     D: Output diagonal factor from the factorization (only used for sparse).
     x: Output array for the solution.
     y: Input right-hand side array.
@@ -2972,6 +2963,132 @@ def factor_solve_i(m, d, M, L, D, x, y):
     _factor_solve_i_dense(m, d, M, x, y, L)
 
 
+@cache_kernel
+def _factor_solve_lu_sparse_fused(nv: int):
+  """Fused sparse LU factorization and solve in a single kernel."""
+
+  @wp.kernel(module="unique", enable_backward=False)
+  def kernel(
+    # Model:
+    D_rownnz: wp.array[int],
+    D_rowadr: wp.array[int],
+    D_diag: wp.array[int],
+    D_colind: wp.array[int],
+    # In:
+    qfrc: wp.array2d[float],
+    # Data out:
+    qacc_out: wp.array2d[float],
+    qLU_out: wp.array3d[float],
+  ):
+    worldid = wp.tid()
+    NV = wp.static(nv)
+
+    # Phase 1: LU factorization (in-place on qLU_out)
+    for i in range(NV):
+      qacc_out[worldid, i] = float(D_rownnz[i])
+
+    # process diagonal elements from n-1 down to 0
+    for r_rev in range(NV):
+      i = NV - 1 - r_rev
+
+      rem_i = int(qacc_out[worldid, i])
+      rowadr_i = D_rowadr[i]
+      ii = rowadr_i + rem_i - 1
+      qacc_out[worldid, i] = float(rem_i - 1)
+
+      # cache diagonal element for row i
+      LUii = qLU_out[worldid, 0, ii]
+
+      # rows j above i (j < i), processed from i-1 down to 0
+      for c in range(i):
+        j = i - 1 - c
+
+        # get address of last remaining element of row j
+        rem_j = int(qacc_out[worldid, j])
+        rowadr_j = D_rowadr[j]
+        ji = rowadr_j + rem_j - 1
+
+        # process row j if (j,i) is non-zero
+        if D_colind[ji] == i:
+          # adjust remaining counter
+          rem_j = rem_j - 1
+          qacc_out[worldid, j] = float(rem_j)
+
+          # (j,i) = (j,i) / (i,i)
+          LUji = qLU_out[worldid, 0, ji] / LUii
+          qLU_out[worldid, 0, ji] = LUji
+
+          # (j,k) = (j,k) - (i,k) * (j,i) for k < i
+          icnt = rowadr_i
+          jcnt = rowadr_j
+          jend = rowadr_j + rem_j
+          while jcnt < jend:
+            col_i = D_colind[icnt]
+            col_j = D_colind[jcnt]
+            if col_i == col_j:
+              qLU_out[worldid, 0, jcnt] = qLU_out[worldid, 0, jcnt] - qLU_out[worldid, 0, icnt] * LUji
+              icnt = icnt + 1
+              jcnt = jcnt + 1
+            elif col_i > col_j:
+              jcnt = jcnt + 1
+            else:
+              icnt = icnt + 1
+
+    # Phase 2: LU solve (backward + forward substitution)
+
+    # Backward substitution: solve (U+I)*qacc = qfrc
+    for k_rev in range(NV):
+      i = NV - 1 - k_rev
+
+      diag_i = D_diag[i]
+      rowadr_i = D_rowadr[i]
+      d1 = diag_i + 1
+      nnz_upper = D_rownnz[i] - d1
+
+      acc = qfrc[worldid, i]
+      for j in range(nnz_upper):
+        adr_j = rowadr_i + d1 + j
+        col = D_colind[adr_j]
+        acc = acc - qLU_out[worldid, 0, adr_j] * qacc_out[worldid, col]
+      qacc_out[worldid, i] = acc
+
+    # Forward substitution: solve L*qacc = qacc
+    for i in range(NV):
+      diag_i = D_diag[i]
+      rowadr_i = D_rowadr[i]
+
+      acc = qacc_out[worldid, i]
+      for j in range(diag_i):
+        adr_j = rowadr_i + j
+        col = D_colind[adr_j]
+        acc = acc - qLU_out[worldid, 0, adr_j] * qacc_out[worldid, col]
+
+      qacc_out[worldid, i] = acc / qLU_out[worldid, 0, rowadr_i + diag_i]
+
+  return kernel
+
+
+@event_scope
+def factor_solve_lu(m: Model, d: Data, qLU: wp.array3d[float], qacc: wp.array2d[float], qfrc: wp.array2d[float]):
+  r"""Factorize and solve non-symmetric implicit system: qacc = A \\ qfrc.
+
+  qLU is overwritten in-place with the LU factors, then used to solve for qacc.
+
+  Args:
+    m: The model containing D-structure sparsity information.
+    d: The data object.
+    qLU: array containing the system matrix, overwritten with LU factors.
+    qacc: output array for the solution.
+    qfrc: input right-hand side.
+  """
+  wp.launch(
+    _factor_solve_lu_sparse_fused(m.nv),
+    dim=(d.nworld,),
+    inputs=[m.D_rownnz, m.D_rowadr, m.D_diag, m.D_colind, qfrc],
+    outputs=[qacc, qLU],
+  )
+
+
 @wp.kernel
 def _subtree_vel_forward(
   # Model:
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py
index 0c804716..1a8ab32b 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py
@@ -13,18 +13,21 @@
 # limitations under the License.
 # ==============================================================================
 
-import dataclasses
 from math import ceil
 from math import sqrt
 
 import warp as wp
 
+from mujoco.mjx.third_party.mujoco_warp._src import island
 from mujoco.mjx.third_party.mujoco_warp._src import math
 from mujoco.mjx.third_party.mujoco_warp._src import smooth
 from mujoco.mjx.third_party.mujoco_warp._src import support
 from mujoco.mjx.third_party.mujoco_warp._src import types
-from mujoco.mjx.third_party.mujoco_warp._src.block_cholesky import create_blocked_cholesky_func
+from mujoco.mjx.third_party.mujoco_warp._src.block_cholesky import create_blocked_cholesky_factorize_solve_func
 from mujoco.mjx.third_party.mujoco_warp._src.block_cholesky import create_blocked_cholesky_solve_func
+from mujoco.mjx.third_party.mujoco_warp._src.types import InverseContext
+from mujoco.mjx.third_party.mujoco_warp._src.types import IslandSolverContext
+from mujoco.mjx.third_party.mujoco_warp._src.types import SolverContext
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import scoped_mathdx_gemm_disabled
@@ -34,55 +37,12 @@ wp.set_module_options({"enable_backward": False})
 _BLOCK_CHOLESKY_DIM = 32
 
 
-@dataclasses.dataclass
-class InverseContext:
-  """Workspace arrays for inverse dynamics."""
-
-  Jaref: wp.array2d[float]
-  search_dot: wp.array[float]
-  gauss: wp.array[float]
-  cost: wp.array[float]
-  prev_cost: wp.array[float]
-  done: wp.array[bool]
-  changed_efc_ids: wp.array2d[int]
-  changed_efc_count: wp.array[int]
-
-
-@dataclasses.dataclass
-class SolverContext:
-  """Workspace arrays for constraint solver."""
-
-  Jaref: wp.array2d[float]
-  search_dot: wp.array[float]
-  gauss: wp.array[float]
-  cost: wp.array[float]
-  prev_cost: wp.array[float]
-  done: wp.array[bool]
-  grad: wp.array2d[float]
-  grad_dot: wp.array[float]
-  Mgrad: wp.array2d[float]
-  search: wp.array2d[float]
-  mv: wp.array2d[float]
-  jv: wp.array2d[float]
-  quad: wp.array2d[wp.vec3]
-  quad_gauss: wp.array[wp.vec3]
-  alpha: wp.array[float]
-  prev_grad: wp.array2d[float]
-  prev_Mgrad: wp.array2d[float]
-  beta: wp.array[float]
-  h: wp.array3d[float]
-  hfactor: wp.array3d[float]
-  # Incremental Hessian update (Newton only)
-  changed_efc_ids: wp.array2d[int]
-  changed_efc_count: wp.array[int]
-
-
 def create_inverse_context(m: types.Model, d: types.Data) -> InverseContext:
   """Create an InverseContext with allocated workspace arrays.
 
   Args:
-    m: Model containing nv, nv_pad, and solver type.
-    d: Data containing nworld and njmax.
+    m: Model.
+    d: Data.
 
   Returns:
     InverseContext with allocated arrays.
@@ -102,12 +62,55 @@ def create_inverse_context(m: types.Model, d: types.Data) -> InverseContext:
   )
 
 
+def create_island_solver_context(m: types.Model, d: types.Data) -> IslandSolverContext:
+  """Create an IslandSolverContext with allocated workspace arrays.
+
+  Args:
+    m: Model.
+    d: Data.
+
+  Returns:
+    IslandSolverContext with allocated arrays.
+  """
+  nworld = d.nworld
+  nv = m.nv
+  nv_pad = m.nv_pad
+  njmax = d.njmax
+  ntree = m.ntree
+
+  alloc_h = m.opt.solver == types.SolverType.NEWTON
+  alloc_island_cg = m.opt.solver == types.SolverType.CG
+
+  return IslandSolverContext(
+    Jaref=wp.empty((nworld, njmax), dtype=float),
+    jv=wp.empty((nworld, njmax), dtype=float),
+    search=wp.empty((nworld, nv), dtype=float),
+    mv=wp.empty((nworld, nv), dtype=float),
+    grad=wp.zeros((nworld, nv_pad), dtype=float),
+    Mgrad=wp.zeros((nworld, nv_pad), dtype=float),
+    prev_grad=wp.empty((nworld, nv), dtype=float) if alloc_island_cg else wp.empty((nworld, 0), dtype=float),
+    prev_Mgrad=wp.empty((nworld, nv), dtype=float) if alloc_island_cg else wp.empty((nworld, 0), dtype=float),
+    h=wp.zeros((nworld, nv_pad, nv_pad), dtype=float) if alloc_h else wp.empty((nworld, 0, 0), dtype=float),
+    # Per-island solver scalars
+    cost=wp.empty((nworld, ntree), dtype=float),
+    prev_cost=wp.empty((nworld, ntree), dtype=float),
+    gauss=wp.empty((nworld, ntree), dtype=float),
+    search_dot=wp.empty((nworld, ntree), dtype=float),
+    grad_dot=wp.empty((nworld, ntree), dtype=float),
+    done=wp.empty((nworld, ntree), dtype=bool),
+    solver_niter=wp.empty((nworld, ntree), dtype=int),
+    beta=wp.empty((nworld, ntree), dtype=float) if alloc_island_cg else wp.empty((nworld, 0), dtype=float),
+    alpha=wp.empty((nworld, ntree), dtype=float),
+    Ma=wp.empty((nworld, nv), dtype=float),
+  )
+
+
 def create_solver_context(m: types.Model, d: types.Data) -> SolverContext:
   """Create a SolverContext with allocated workspace arrays.
 
   Args:
-    m: Model containing nv, nv_pad, and solver type.
-    d: Data containing nworld and njmax.
+    m: Model.
+    d: Data.
 
   Returns:
     SolverContext with allocated arrays.
@@ -1669,7 +1672,7 @@ def _linesearch(m: types.Model, d: types.Data, ctx: SolverContext, cost: wp.arra
     ctx: SolverContext
     cost: Scratch array for storing costs per (world, alpha) - used for parallel mode
   """
-  # mv = qM @ search (common to both parallel and iterative)
+  # mv = M @ search (common to both parallel and iterative)
   support.mul_m(m, d, ctx.mv, ctx.search, skip=ctx.done)
 
   # Fuse jv computation in-kernel for small nv (iterative only, dense only)
@@ -2063,9 +2066,9 @@ def update_gradient_h_incremental(
   # Out:
   ctx_h_out: wp.array3d[float],
 ):
-  """Incrementally update lower triangle of H for changed constraints.
+  """Incrementally update upper triangle of H for changed constraints.
 
-  Each thread handles one (i, j) element of the lower triangle.
+  Each thread handles one unique (i, j) element and writes it to the upper triangle.
   For each changed constraint, adds or subtracts D * J[i] * J[j].
   """
   worldid, elementid = wp.tid()
@@ -2074,28 +2077,28 @@ def update_gradient_h_incremental(
   if n_changes == 0:
     return
 
-  # Lower triangle index: elementid -> (i, j) where i >= j
-  i = (int(wp.sqrt(float(1 + 8 * elementid))) - 1) // 2
-  j = elementid - (i * (i + 1)) // 2
+  # Upper-triangle enumeration: elementid -> (row, col) where row <= col.
+  col = (int(wp.sqrt(float(1 + 8 * elementid))) - 1) // 2
+  row = elementid - (col * (col + 1)) // 2
 
   delta = float(0.0)
   for change_idx in range(n_changes):
     efcid = changed_ids_in[worldid, change_idx]
-    Ji = efc_J_in[worldid, efcid, i]
-    if Ji == 0.0:
+    Jrow = efc_J_in[worldid, efcid, row]
+    if Jrow == 0.0:
       continue
-    Jj = efc_J_in[worldid, efcid, j]
-    if Jj == 0.0:
+    Jcol = efc_J_in[worldid, efcid, col]
+    if Jcol == 0.0:
       continue
 
     D = efc_D_in[worldid, efcid]
     if efc_state_in[worldid, efcid] == types.ConstraintState.QUADRATIC.value:
-      delta += D * Ji * Jj
+      delta += D * Jrow * Jcol
     else:
-      delta -= D * Ji * Jj
+      delta -= D * Jrow * Jcol
 
   if delta != 0.0:
-    ctx_h_out[worldid, i, j] += delta
+    ctx_h_out[worldid, row, col] += delta
 
 
 @wp.kernel
@@ -2113,7 +2116,7 @@ def update_gradient_h_incremental_sparse(
   # Out:
   ctx_h_out: wp.array3d[float],
 ):
-  """Incrementally update lower triangle of H for changed constraints (sparse J)."""
+  """Incrementally update upper triangle of H for changed constraints (sparse J)."""
   worldid, change_idx = wp.tid()
 
   n_changes = changed_count_in[worldid]
@@ -2144,8 +2147,8 @@ def update_gradient_h_incremental_sparse(
         continue
       colindj = efc_J_colind_in[worldid, 0, sparseidj]
       h = sign * Ji * Jj
-      # Ensure lower triangle: larger index first
-      if colindi >= colindj:
+      # Ensure upper triangle: smaller index first.
+      if colindi <= colindj:
         wp.atomic_add(ctx_h_out[worldid, colindi], colindj, h)
       else:
         wp.atomic_add(ctx_h_out[worldid, colindj], colindi, h)
@@ -2257,12 +2260,13 @@ def update_gradient_grad(
 
 
 @wp.kernel
-def update_gradient_set_h_qM_lower_sparse(
+def update_gradient_set_h_M_upper_sparse(
   # Model:
-  qM_fullm_i: wp.array[int],
-  qM_fullm_j: wp.array[int],
+  M_fullm_upper_i: wp.array[int],
+  M_fullm_upper_j: wp.array[int],
+  M_fullm_upper_elemid: wp.array[int],
   # Data in:
-  qM_in: wp.array3d[float],
+  M_in: wp.array3d[float],
   # In:
   ctx_done_in: wp.array[bool],
   # Out:
@@ -2273,9 +2277,10 @@ def update_gradient_set_h_qM_lower_sparse(
   if ctx_done_in[worldid]:
     return
 
-  i = qM_fullm_i[elementid]
-  j = qM_fullm_j[elementid]
-  ctx_h_out[worldid, i, j] += qM_in[worldid, 0, elementid]
+  i = M_fullm_upper_i[elementid]
+  j = M_fullm_upper_j[elementid]
+  madr = M_fullm_upper_elemid[elementid]
+  ctx_h_out[worldid, i, j] += M_in[worldid, 0, madr]
 
 
 @wp.func
@@ -2317,12 +2322,12 @@ def update_gradient_JTDAJ_sparse_tiled(tile_size: int, njmax: int):
 
     nefc = nefc_in[worldid]
 
-    # get lower diagonal index
-    i = (int(sqrt(float(1 + 8 * elementid))) - 1) // 2
-    j = elementid - (i * (i + 1)) // 2
+    # Upper-triangle tile index: elementid -> (row, col) where row <= col.
+    col = (int(sqrt(float(1 + 8 * elementid))) - 1) // 2
+    row = elementid - (col * (col + 1)) // 2
 
-    offset_i = i * TILE_SIZE
-    offset_j = j * TILE_SIZE
+    offset_row = row * TILE_SIZE
+    offset_col = col * TILE_SIZE
 
     sum_val = wp.tile_zeros(shape=(TILE_SIZE, TILE_SIZE), dtype=wp.float32)
 
@@ -2334,12 +2339,12 @@ def update_gradient_JTDAJ_sparse_tiled(tile_size: int, njmax: int):
       # AD: leaving bounds-check disabled here because I'm not entirely sure that
       # everything always hits the fast path. The padding takes care of any
       # potential OOB accesses.
-      J_ki = wp.tile_load(efc_J_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(k, offset_i), bounds_check=False)
+      J_krow = wp.tile_load(efc_J_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(k, offset_row), bounds_check=False)
 
-      if offset_i != offset_j:
-        J_kj = wp.tile_load(efc_J_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(k, offset_j), bounds_check=False)
+      if offset_row != offset_col:
+        J_kcol = wp.tile_load(efc_J_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(k, offset_col), bounds_check=False)
       else:
-        wp.tile_assign(J_kj, J_ki, (0, 0))
+        wp.tile_assign(J_kcol, J_krow, (0, 0))
 
       D_k = wp.tile_load(efc_D_in[worldid], shape=TILE_SIZE, offset=k, bounds_check=False)
       state = wp.tile_load(efc_state_in[worldid], shape=TILE_SIZE, offset=k, bounds_check=False)
@@ -2353,13 +2358,13 @@ def update_gradient_JTDAJ_sparse_tiled(tile_size: int, njmax: int):
       active_tile = wp.tile_map(active_check, tid_tile, threshold_tile)
       D_k = wp.tile_map(wp.mul, active_tile, D_k)
 
-      J_ki = wp.tile_map(wp.mul, wp.tile_transpose(J_ki), wp.tile_broadcast(D_k, shape=(TILE_SIZE, TILE_SIZE)))
+      J_krow = wp.tile_map(wp.mul, wp.tile_transpose(J_krow), wp.tile_broadcast(D_k, shape=(TILE_SIZE, TILE_SIZE)))
 
-      sum_val += wp.tile_matmul(J_ki, J_kj)
+      sum_val += wp.tile_matmul(J_krow, J_kcol)
 
     # AD: setting bounds_check to True explicitly here because for some reason it was
     # slower to disable it.
-    wp.tile_store(ctx_h_out[worldid], sum_val, offset=(offset_i, offset_j), bounds_check=True)
+    wp.tile_store(ctx_h_out[worldid], sum_val, offset=(offset_row, offset_col), bounds_check=True)
 
   return kernel
 
@@ -2375,7 +2380,7 @@ def update_gradient_JTDAJ_dense_tiled(nv_pad: int, tile_size: int, njmax: int):
   def kernel(
     # Data in:
     nefc_in: wp.array[int],
-    qM_in: wp.array3d[float],
+    M_in: wp.array3d[float],
     efc_J_in: wp.array3d[float],
     efc_D_in: wp.array2d[float],
     efc_state_in: wp.array2d[int],
@@ -2391,7 +2396,7 @@ def update_gradient_JTDAJ_dense_tiled(nv_pad: int, tile_size: int, njmax: int):
 
     nefc = nefc_in[worldid]
 
-    sum_val = wp.tile_load(qM_in[worldid], shape=(nv_pad, nv_pad), bounds_check=True)
+    sum_val = wp.tile_load(M_in[worldid], shape=(nv_pad, nv_pad), bounds_check=True)
 
     # Each tile processes one output tile by looping over all constraints
     for k in range(0, njmax, TILE_SIZE_K):
@@ -2746,9 +2751,9 @@ def update_gradient_cholesky(tile_size: int):
       return
 
     mat_tile = wp.tile_load(h_in[worldid], shape=(TILE_SIZE, TILE_SIZE))
-    fact_tile = wp.tile_cholesky(mat_tile)
+    fact_tile = wp.tile_cholesky(mat_tile, fill_mode="upper")
     input_tile = wp.tile_load(ctx_grad_in[worldid], shape=TILE_SIZE)
-    output_tile = wp.tile_cholesky_solve(fact_tile, input_tile)
+    output_tile = wp.tile_cholesky_solve(fact_tile, input_tile, fill_mode="upper")
     wp.tile_store(ctx_Mgrad_out[worldid], output_tile)
 
   return kernel
@@ -2777,14 +2782,14 @@ def update_gradient_cholesky_blocked(tile_size: int, matrix_size: int):
     # runtime input is needed for the loop bounds, otherwise warp will unroll
     # unconditionally leading to shared memory capacity issues.
 
-    wp.static(create_blocked_cholesky_func(TILE_SIZE))(ctx_h_in[worldid], matrix_size, ctx_hfactor[worldid])
-    wp.static(create_blocked_cholesky_solve_func(TILE_SIZE, matrix_size))(
-      ctx_hfactor[worldid], ctx_grad_in[worldid], matrix_size, ctx_Mgrad_out[worldid]
+    wp.static(create_blocked_cholesky_factorize_solve_func(TILE_SIZE, matrix_size))(
+      ctx_h_in[worldid], ctx_grad_in[worldid], matrix_size, ctx_hfactor[worldid], ctx_Mgrad_out[worldid]
     )
 
   return kernel
 
 
+@cache_kernel
 def update_gradient_cholesky_blocked_skip_unchanged(tile_size: int, matrix_size: int):
   """Blocked Cholesky that skips factorization when no constraints changed."""
 
@@ -2806,11 +2811,13 @@ def update_gradient_cholesky_blocked_skip_unchanged(tile_size: int, matrix_size:
       return
 
     if changed_count_in[worldid] > 0:
-      wp.static(create_blocked_cholesky_func(TILE_SIZE))(ctx_h_in[worldid], matrix_size, ctx_hfactor[worldid])
-
-    wp.static(create_blocked_cholesky_solve_func(TILE_SIZE, matrix_size))(
-      ctx_hfactor[worldid], ctx_grad_in[worldid], matrix_size, ctx_Mgrad_out[worldid]
-    )
+      wp.static(create_blocked_cholesky_factorize_solve_func(TILE_SIZE, matrix_size))(
+        ctx_h_in[worldid], ctx_grad_in[worldid], matrix_size, ctx_hfactor[worldid], ctx_Mgrad_out[worldid]
+      )
+    else:
+      wp.static(create_blocked_cholesky_solve_func(TILE_SIZE, matrix_size))(
+        ctx_hfactor[worldid], ctx_grad_in[worldid], matrix_size, ctx_Mgrad_out[worldid]
+      )
 
   return kernel
 
@@ -2913,9 +2920,9 @@ def _JTDAJ_sparse(
         colindj = efc_J_colind_in[worldid, 0, sparseidj]
 
       h = Ji * Jj * efc_D
-      # Store in lower triangle only: ensure row >= col
-      row = wp.max(colindi, colindj)
-      col = wp.min(colindi, colindj)
+      # Store in upper triangle only: ensure row <= col.
+      row = wp.min(colindi, colindj)
+      col = wp.max(colindi, colindj)
       wp.atomic_add(h_out[worldid, row], col, h)
 
 
@@ -2933,7 +2940,7 @@ def _update_gradient(m: types.Model, d: types.Data, ctx: SolverContext):
   if m.opt.solver == types.SolverType.CG:
     smooth.solve_m(m, d, ctx.Mgrad, ctx.grad)
   elif m.opt.solver == types.SolverType.NEWTON:
-    # h = qM + (efc_J.T * efc_D * active) @ efc_J
+    # h = M + (efc_J.T * efc_D * active) @ efc_J
     if m.is_sparse:
       ctx.h.zero_()
       wp.launch(
@@ -2944,9 +2951,9 @@ def _update_gradient(m: types.Model, d: types.Data, ctx: SolverContext):
       )
 
       wp.launch(
-        update_gradient_set_h_qM_lower_sparse,
-        dim=(d.nworld, m.qM_fullm_i.size),
-        inputs=[m.qM_fullm_i, m.qM_fullm_j, d.qM, ctx.done],
+        update_gradient_set_h_M_upper_sparse,
+        dim=(d.nworld, m.M_fullm_upper_i.size),
+        inputs=[m.M_fullm_upper_i, m.M_fullm_upper_j, m.M_fullm_upper_elemid, d.M, ctx.done],
         outputs=[ctx.h],
       )
     else:
@@ -2956,7 +2963,7 @@ def _update_gradient(m: types.Model, d: types.Data, ctx: SolverContext):
           dim=d.nworld,
           inputs=[
             d.nefc,
-            d.qM,
+            d.M,
             d.efc.J,
             d.efc.D,
             d.efc.state,
@@ -3066,7 +3073,7 @@ def _update_gradient_incremental(m: types.Model, d: types.Data, ctx: SolverConte
     outputs=[ctx.grad, ctx.grad_dot],
   )
 
-  # Update lower triangle of H with delta from changed constraints
+  # Update upper triangle of H with delta from changed constraints.
   if m.is_sparse:
     wp.launch(
       update_gradient_h_incremental_sparse,
@@ -3084,10 +3091,10 @@ def _update_gradient_incremental(m: types.Model, d: types.Data, ctx: SolverConte
       outputs=[ctx.h],
     )
   else:
-    lower_tri_dim = m.nv * (m.nv + 1) // 2
+    tri_dim = m.nv * (m.nv + 1) // 2
     wp.launch(
       update_gradient_h_incremental,
-      dim=(d.nworld, lower_tri_dim),
+      dim=(d.nworld, tri_dim),
       inputs=[
         d.efc.J,
         d.efc.D,
@@ -3327,7 +3334,7 @@ def init_context(m: types.Model, d: types.Data, ctx: SolverContext | InverseCont
     outputs=[ctx.Jaref],
   )
 
-  # Ma = qM @ qacc
+  # Ma = M @ qacc
   support.mul_m(m, d, d.efc.Ma, d.qacc, skip=ctx.done)
 
   _update_constraint(m, d, ctx)
@@ -3342,8 +3349,17 @@ def solve(m: types.Model, d: types.Data):
     wp.copy(d.qacc, d.qacc_smooth)
     d.solver_niter.fill_(0)
   else:
-    ctx = create_solver_context(m, d)
-    _solve(m, d, ctx)
+    if m.ntree > 1 and not (m.opt.disableflags & types.DisableBit.ISLAND):
+      ctx = create_island_solver_context(m, d)
+      island.compute_island_mapping(m, d, ctx)
+      island.gather_island_inputs(m, d, ctx)
+      _solve_islands(m, d, ctx)
+      # Ma is needed by Euler/implicit integrators for implicit damping
+      scatter_Ma = m.opt.integrator != types.IntegratorType.RK4
+      island.scatter_island_results(m, d, ctx, scatter_Ma=scatter_Ma)
+    else:
+      ctx = create_solver_context(m, d)
+      _solve(m, d, ctx)
 
 
 def _solve(m: types.Model, d: types.Data, ctx: SolverContext):
@@ -3384,3 +3400,2104 @@ def _solve(m: types.Model, d: types.Data, ctx: SolverContext):
     # It should be removed when JAX becomes compatible.
     for _ in range(m.opt.iterations):
       _solver_iteration(m, d, ctx, step_size_cost, nsolving)
+
+
+# TODO(team): Consolidate monolithic and island solver code where possible
+@event_scope
+def _solve_islands(m: types.Model, d: types.Data, ctx: IslandSolverContext):
+  """Solve constraints for all islands in parallel.
+
+  All islands are processed simultaneously. Island-local arrays in ctx
+  (iacc, iefc_J, iefc_D, etc.) are indexed by idof/iefc, and each thread
+  determines its island via idof_islandid/iefc_islandid lookup tables.
+  """
+  # Initialize iacc from warmstart or smooth
+  if not (m.opt.disableflags & types.DisableBit.WARMSTART):
+    wp.launch(
+      gather_warmstart_island,
+      dim=(d.nworld, m.nv),
+      inputs=[d.nidof, d.qacc_warmstart, d.map_idof2dof],
+      outputs=[d.iqacc],
+    )
+  else:
+    wp.copy(d.iqacc, d.iqacc_smooth)
+
+  # Initialize island context
+  init_context_island(m, d, ctx)
+
+  # search = -Mgrad
+  wp.launch(
+    solve_init_search_island,
+    dim=(d.nworld, m.nv),
+    inputs=[d.nidof, ctx.Mgrad, d.dof_islandid, ctx.done],
+    outputs=[ctx.search, ctx.search_dot],
+  )
+
+  # nsolving tracks how many active islands still have unconverged globally
+  nsolving = wp.zeros((1,), dtype=int)
+  wp.launch(
+    solve_init_nsolving_island,
+    dim=d.nworld,
+    inputs=[d.nisland],
+    outputs=[nsolving],
+  )
+
+  if m.opt.iterations != 0 and m.opt.graph_conditional:
+    wp.capture_while(
+      nsolving,
+      while_body=_solver_iteration_island,
+      m=m,
+      d=d,
+      ctx=ctx,
+      nsolving=nsolving,
+    )
+  else:
+    for _ in range(m.opt.iterations):
+      _solver_iteration_island(m, d, ctx, nsolving)
+
+
+@wp.kernel
+def gather_warmstart_island(
+  # Data in:
+  nidof_in: wp.array[int],
+  qacc_warmstart_in: wp.array2d[float],
+  map_idof2dof_in: wp.array2d[int],
+  # Out:
+  iacc_out: wp.array2d[float],
+):
+  """Gather qacc_warmstart into island-local order."""
+  worldid, idofid = wp.tid()
+
+  if idofid >= nidof_in[worldid]:
+    return
+
+  dof = map_idof2dof_in[worldid, idofid]
+  iacc_out[worldid, idofid] = qacc_warmstart_in[worldid, dof]
+
+
+@wp.kernel
+def solve_init_efc_island(
+  # Data in:
+  nisland_in: wp.array[int],
+  # Out:
+  island_cost_out: wp.array2d[float],
+  island_search_dot_out: wp.array2d[float],
+  island_done_out: wp.array2d[bool],
+  island_solver_niter_out: wp.array2d[int],
+):
+  """Initialize per-island solver scalars."""
+  worldid, islandid = wp.tid()
+
+  if islandid >= nisland_in[worldid]:
+    return
+
+  island_cost_out[worldid, islandid] = 0.0
+  island_search_dot_out[worldid, islandid] = 0.0
+  island_done_out[worldid, islandid] = False
+  island_solver_niter_out[worldid, islandid] = 0
+
+
+@wp.kernel
+def solve_init_jaref_island(
+  # Model:
+  is_sparse: bool,
+  # Data in:
+  nefc_in: wp.array[int],
+  island_nv_in: wp.array2d[int],
+  njmax_in: int,
+  # In:
+  island_idofadr_in: wp.array2d[int],
+  iefc_J_rownnz_in: wp.array2d[int],
+  iefc_J_rowadr_in: wp.array2d[int],
+  iefc_J_colind_in: wp.array3d[int],
+  iefc_J_in: wp.array3d[float],
+  iacc_in: wp.array2d[float],
+  iefc_aref_in: wp.array2d[float],
+  iefc_islandid_in: wp.array2d[int],
+  island_done_in: wp.array2d[bool],
+  # Out:
+  Jaref_out: wp.array2d[float],
+):
+  """Jaref[iefcid] = iefc_J[iefcid] @ iacc - iefc_aref[iefcid] for all island EFCs."""
+  worldid, iefcid = wp.tid()
+
+  if iefcid >= wp.min(njmax_in, nefc_in[worldid]):
+    return
+
+  islandid = iefc_islandid_in[worldid, iefcid]
+  if islandid < 0:
+    return
+  if island_done_in[worldid, islandid]:
+    return
+
+  acc = float(0.0)
+  if is_sparse:
+    rownnz = iefc_J_rownnz_in[worldid, iefcid]
+    rowadr = iefc_J_rowadr_in[worldid, iefcid]
+    for k in range(rownnz):
+      adr = rowadr + k
+      Ji = iefc_J_in[worldid, 0, adr]
+      idof = iefc_J_colind_in[worldid, 0, adr]
+      acc += Ji * iacc_in[worldid, idof]
+  else:
+    idofadr = island_idofadr_in[worldid, islandid]
+    inv = island_nv_in[worldid, islandid]
+    for i in range(inv):
+      idof = idofadr + i
+      acc += iefc_J_in[worldid, iefcid, idof] * iacc_in[worldid, idof]
+
+  Jaref_out[worldid, iefcid] = acc - iefc_aref_in[worldid, iefcid]
+
+
+@wp.kernel
+def solve_init_search_island(
+  # Data in:
+  nidof_in: wp.array[int],
+  # In:
+  Mgrad_in: wp.array2d[float],
+  idof_islandid_in: wp.array2d[int],
+  island_done_in: wp.array2d[bool],
+  # Out:
+  search_out: wp.array2d[float],
+  island_search_dot_out: wp.array2d[float],
+):
+  """Search = -Mgrad for all island DOFs."""
+  worldid, idofid = wp.tid()
+
+  if idofid >= nidof_in[worldid]:
+    return
+
+  islandid = idof_islandid_in[worldid, idofid]
+  if islandid < 0:
+    return
+  if island_done_in[worldid, islandid]:
+    return
+
+  s = -Mgrad_in[worldid, idofid]
+  search_out[worldid, idofid] = s
+  wp.atomic_add(island_search_dot_out, worldid, islandid, s * s)
+
+
+@wp.kernel
+def update_constraint_init_cost_island(
+  # Data in:
+  nisland_in: wp.array[int],
+  # In:
+  island_cost_in: wp.array2d[float],
+  island_done_in: wp.array2d[bool],
+  # Out:
+  island_gauss_out: wp.array2d[float],
+  island_cost_out: wp.array2d[float],
+  island_prev_cost_out: wp.array2d[float],
+):
+  """Save prev_cost and zero cost/gauss per island."""
+  worldid, islandid = wp.tid()
+
+  if islandid >= nisland_in[worldid]:
+    return
+  if island_done_in[worldid, islandid]:
+    return
+
+  island_prev_cost_out[worldid, islandid] = island_cost_in[worldid, islandid]
+  island_cost_out[worldid, islandid] = 0.0
+  island_gauss_out[worldid, islandid] = 0.0
+
+
+@wp.kernel
+def update_constraint_efc_island(
+  # Model:
+  opt_impratio_invsqrt: wp.array[float],
+  # Data in:
+  nefc_in: wp.array[int],
+  contact_friction_in: wp.array[types.vec5],
+  contact_dim_in: wp.array[int],
+  contact_efc_address_in: wp.array2d[int],
+  island_nefc_in: wp.array2d[int],
+  island_ne_in: wp.array2d[int],
+  island_nf_in: wp.array2d[int],
+  island_efcadr_in: wp.array2d[int],
+  map_efc2iefc_in: wp.array2d[int],
+  njmax_in: int,
+  nacon_in: wp.array[int],
+  # In:
+  iefc_type_in: wp.array2d[int],
+  iefc_id_in: wp.array2d[int],
+  iefc_D_in: wp.array2d[float],
+  iefc_frictionloss_in: wp.array2d[float],
+  Jaref_in: wp.array2d[float],
+  iefc_islandid_in: wp.array2d[int],
+  island_done_in: wp.array2d[bool],
+  # Out:
+  iefc_force_out: wp.array2d[float],
+  iefc_state_out: wp.array2d[int],
+  island_cost_out: wp.array2d[float],
+):
+  """Compute force, state, and cost for each island constraint."""
+  worldid, iefcid = wp.tid()
+
+  if iefcid >= wp.min(njmax_in, nefc_in[worldid]):
+    return
+
+  islandid = iefc_islandid_in[worldid, iefcid]
+  if islandid < 0:
+    return
+  if island_done_in[worldid, islandid]:
+    return
+
+  # Local position within island
+  iefcadr = island_efcadr_in[worldid, islandid]
+  local_iefcid = iefcid - iefcadr
+  ine = island_ne_in[worldid, islandid]
+  inf = island_nf_in[worldid, islandid]
+
+  jaref = Jaref_in[worldid, iefcid]
+  D = iefc_D_in[worldid, iefcid]
+
+  force = float(0.0)
+  state = types.ConstraintState.SATISFIED.value
+  cost = float(0.0)
+
+  # Determine constraint category by position within island
+  if local_iefcid < ine:
+    # Equality constraint: always quadratic
+    force = -D * jaref
+    state = types.ConstraintState.QUADRATIC.value
+    cost = 0.5 * D * jaref * jaref
+  elif local_iefcid < ine + inf:
+    # Friction constraint
+    f = iefc_frictionloss_in[worldid, iefcid]
+    rf = math.safe_div(f, D)
+    if jaref <= -rf:
+      force = f
+      state = types.ConstraintState.LINEARNEG.value
+      cost = -f * (0.5 * rf + jaref)
+    elif jaref >= rf:
+      force = -f
+      state = types.ConstraintState.LINEARPOS.value
+      cost = -f * (0.5 * rf - jaref)
+    else:
+      force = -D * jaref
+      state = types.ConstraintState.QUADRATIC.value
+      cost = 0.5 * D * jaref * jaref
+  elif iefc_type_in[worldid, iefcid] != types.ConstraintType.CONTACT_ELLIPTIC:
+    # Limit, frictionless contact, pyramidal friction cone contact
+    if jaref >= 0.0:
+      force = 0.0
+      state = types.ConstraintState.SATISFIED.value
+      cost = 0.0
+    else:
+      force = -D * jaref
+      state = types.ConstraintState.QUADRATIC.value
+      cost = 0.5 * D * jaref * jaref
+  else:
+    # Elliptic friction cone contact
+    conid = iefc_id_in[worldid, iefcid]
+
+    if conid >= nacon_in[0]:
+      return
+
+    dim = contact_dim_in[conid]
+    friction = contact_friction_in[conid]
+    mu = friction[0] * opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]]
+
+    # Get island-local index for lead row
+    efcid0_global = contact_efc_address_in[conid, 0]
+    if efcid0_global < 0:
+      return
+    ic0 = map_efc2iefc_in[worldid, efcid0_global]
+
+    N = Jaref_in[worldid, ic0] * mu
+
+    ufrictionj = float(0.0)
+    TT = float(0.0)
+    for j in range(1, dim):
+      efcidj_global = contact_efc_address_in[conid, j]
+      if efcidj_global < 0:
+        return
+      icj = map_efc2iefc_in[worldid, efcidj_global]
+      frictionj = friction[j - 1]
+      uj = Jaref_in[worldid, icj] * frictionj
+      TT += uj * uj
+      if iefcid == icj:
+        ufrictionj = uj * frictionj
+
+    if TT <= 0.0:
+      T = 0.0
+    else:
+      T = wp.sqrt(TT)
+
+    # top zone
+    if (N >= mu * T) or ((T <= 0.0) and (N >= 0.0)):
+      force = 0.0
+      state = types.ConstraintState.SATISFIED.value
+    # bottom zone
+    elif (mu * N + T <= 0.0) or ((T <= 0.0) and (N < 0.0)):
+      force = -D * jaref
+      state = types.ConstraintState.QUADRATIC.value
+      cost = 0.5 * D * jaref * jaref
+    # middle zone
+    else:
+      dm = math.safe_div(iefc_D_in[worldid, ic0], mu * mu * (1.0 + mu * mu))
+      nmt = N - mu * T
+
+      force_normal = -dm * nmt * mu
+
+      if iefcid == ic0:
+        force = force_normal
+        cost = 0.5 * dm * nmt * nmt
+      else:
+        force = -math.safe_div(force_normal, T) * ufrictionj
+
+      state = types.ConstraintState.CONE.value
+
+  iefc_force_out[worldid, iefcid] = force
+  iefc_state_out[worldid, iefcid] = state
+  wp.atomic_add(island_cost_out, worldid, islandid, cost)
+
+
+@wp.kernel
+def update_constraint_init_qfrc_constraint_dense_island(
+  # Data in:
+  nefc_in: wp.array[int],
+  nidof_in: wp.array[int],
+  island_nefc_in: wp.array2d[int],
+  island_efcadr_in: wp.array2d[int],
+  njmax_in: int,
+  # In:
+  iefc_J_in: wp.array3d[float],
+  iefc_force_in: wp.array2d[float],
+  idof_islandid_in: wp.array2d[int],
+  island_done_in: wp.array2d[bool],
+  # Out:
+  ifrc_constraint_out: wp.array2d[float],
+):
+  """ifrc_constraint = iefc_J.T @ iefc_force for all island DOFs."""
+  worldid, idofid = wp.tid()
+
+  if idofid >= nidof_in[worldid]:
+    ifrc_constraint_out[worldid, idofid] = 0.0
+    return
+
+  islandid = idof_islandid_in[worldid, idofid]
+  if islandid < 0:
+    ifrc_constraint_out[worldid, idofid] = 0.0
+    return
+  if island_done_in[worldid, islandid]:
+    return
+
+  iefcadr = island_efcadr_in[worldid, islandid]
+  inefc = island_nefc_in[worldid, islandid]
+  acc = float(0.0)
+  for iefcid in range(iefcadr, iefcadr + inefc):
+    acc += iefc_J_in[worldid, iefcid, idofid] * iefc_force_in[worldid, iefcid]
+
+  ifrc_constraint_out[worldid, idofid] = acc
+
+
+@wp.kernel
+def update_constraint_init_qfrc_constraint_sparse_island(
+  # Data in:
+  nefc_in: wp.array[int],
+  njmax_in: int,
+  # In:
+  iefc_J_rownnz_in: wp.array2d[int],
+  iefc_J_rowadr_in: wp.array2d[int],
+  iefc_J_colind_in: wp.array3d[int],
+  iefc_J_in: wp.array3d[float],
+  iefc_force_in: wp.array2d[float],
+  iefc_islandid_in: wp.array2d[int],
+  island_done_in: wp.array2d[bool],
+  # Out:
+  ifrc_constraint_out: wp.array2d[float],
+):
+  """ifrc_constraint += iefc_J.T @ iefc_force for all island EFCs (sparse parallel per EFC)."""
+  worldid, iefcid = wp.tid()
+
+  if iefcid >= wp.min(njmax_in, nefc_in[worldid]):
+    return
+
+  islandid = iefc_islandid_in[worldid, iefcid]
+  if islandid < 0:
+    return
+
+  rownnz = iefc_J_rownnz_in[worldid, iefcid]
+  rowadr = iefc_J_rowadr_in[worldid, iefcid]
+  force = iefc_force_in[worldid, iefcid]
+
+  for k in range(rownnz):
+    adr = rowadr + k
+    Ji = iefc_J_in[worldid, 0, adr]
+    idof = iefc_J_colind_in[worldid, 0, adr]
+    wp.atomic_add(ifrc_constraint_out, worldid, idof, Ji * force)
+
+
+@wp.kernel
+def update_constraint_gauss_cost_island(
+  # Data in:
+  nidof_in: wp.array[int],
+  # In:
+  iacc_in: wp.array2d[float],
+  ifrc_smooth_in: wp.array2d[float],
+  iacc_smooth_in: wp.array2d[float],
+  iMa_in: wp.array2d[float],
+  idof_islandid_in: wp.array2d[int],
+  island_done_in: wp.array2d[bool],
+  # Out:
+  island_gauss_out: wp.array2d[float],
+  island_cost_out: wp.array2d[float],
+):
+  """Gauss cost: 0.5 * (Ma - qfrc_smooth).T @ (qacc - qacc_smooth) per island."""
+  worldid, idofid = wp.tid()
+
+  if idofid >= nidof_in[worldid]:
+    return
+
+  islandid = idof_islandid_in[worldid, idofid]
+  if islandid < 0:
+    return
+  if island_done_in[worldid, islandid]:
+    return
+
+  dq = iacc_in[worldid, idofid] - iacc_smooth_in[worldid, idofid]
+  df = iMa_in[worldid, idofid] - ifrc_smooth_in[worldid, idofid]
+  gauss = 0.5 * df * dq
+
+  wp.atomic_add(island_gauss_out, worldid, islandid, gauss)
+  wp.atomic_add(island_cost_out, worldid, islandid, gauss)
+
+
+@wp.kernel
+def update_gradient_grad_island(
+  # Data in:
+  nidof_in: wp.array[int],
+  # In:
+  ifrc_smooth_in: wp.array2d[float],
+  ifrc_constraint_in: wp.array2d[float],
+  iMa_in: wp.array2d[float],
+  idof_islandid_in: wp.array2d[int],
+  island_done_in: wp.array2d[bool],
+  # Out:
+  grad_out: wp.array2d[float],
+  island_grad_dot_out: wp.array2d[float],
+):
+  """Grad = Ma - qfrc_smooth - qfrc_constraint, grad_dot per island."""
+  worldid, idofid = wp.tid()
+
+  if idofid >= nidof_in[worldid]:
+    return
+
+  islandid = idof_islandid_in[worldid, idofid]
+  if islandid < 0:
+    return
+  if island_done_in[worldid, islandid]:
+    return
+
+  g = iMa_in[worldid, idofid] - ifrc_smooth_in[worldid, idofid] - ifrc_constraint_in[worldid, idofid]
+  grad_out[worldid, idofid] = g
+  wp.atomic_add(island_grad_dot_out, worldid, islandid, g * g)
+
+
+@wp.kernel
+def linesearch_jv_island(
+  # Model:
+  is_sparse: bool,
+  # Data in:
+  nefc_in: wp.array[int],
+  nidof_in: wp.array[int],
+  island_nv_in: wp.array2d[int],
+  njmax_in: int,
+  # In:
+  island_idofadr_in: wp.array2d[int],
+  iefc_J_rownnz_in: wp.array2d[int],
+  iefc_J_rowadr_in: wp.array2d[int],
+  iefc_J_colind_in: wp.array3d[int],
+  iefc_J_in: wp.array3d[float],
+  search_in: wp.array2d[float],
+  iefc_islandid_in: wp.array2d[int],
+  island_done_in: wp.array2d[bool],
+  # Out:
+  jv_out: wp.array2d[float],
+):
+  """Jv = iefc_J @ search for all island EFCs."""
+  worldid, iefcid = wp.tid()
+
+  if iefcid >= wp.min(njmax_in, nefc_in[worldid]):
+    return
+
+  islandid = iefc_islandid_in[worldid, iefcid]
+  if islandid < 0:
+    return
+  if island_done_in[worldid, islandid]:
+    return
+
+  acc = float(0.0)
+  if is_sparse:
+    rownnz = iefc_J_rownnz_in[worldid, iefcid]
+    rowadr = iefc_J_rowadr_in[worldid, iefcid]
+    for k in range(rownnz):
+      adr = rowadr + k
+      Ji = iefc_J_in[worldid, 0, adr]
+      idof = iefc_J_colind_in[worldid, 0, adr]
+      acc += Ji * search_in[worldid, idof]
+  else:
+    idofadr = island_idofadr_in[worldid, islandid]
+    inv = island_nv_in[worldid, islandid]
+    for i in range(inv):
+      idof = idofadr + i
+      acc += iefc_J_in[worldid, iefcid, idof] * search_in[worldid, idof]
+
+  jv_out[worldid, iefcid] = acc
+
+
+@wp.func
+def _eval_elliptic_cost_island(
+  # Model:
+  opt_impratio_invsqrt: float,  # kernel_analyzer: off
+  # Data in:
+  contact_friction_in: wp.array[types.vec5],
+  contact_dim_in: wp.array[int],
+  contact_efc_address_in: wp.array2d[int],
+  map_efc2iefc_in: wp.array2d[int],
+  # In:
+  alpha: float,
+  conid: int,
+  iefc_D_in: wp.array2d[float],
+  Jaref_in: wp.array2d[float],
+  jv_in: wp.array2d[float],
+  worldid: int,
+) -> wp.vec3:
+  dim = contact_dim_in[conid]
+  friction = contact_friction_in[conid]
+  mu = friction[0] * opt_impratio_invsqrt
+
+  ic0 = map_efc2iefc_in[worldid, contact_efc_address_in[conid, 0]]
+  D0 = iefc_D_in[worldid, ic0]
+  ja0 = Jaref_in[worldid, ic0]
+  jv0 = jv_in[worldid, ic0]
+
+  # Bottom-zone quad for the full contact (scalar quadratic over all rows)
+  quad = wp.vec3(0.5 * ja0 * ja0 * D0, jv0 * ja0 * D0, 0.5 * jv0 * jv0 * D0)
+
+  u0 = ja0 * mu
+  v0 = jv0 * mu
+  uu = float(0.0)
+  uv = float(0.0)
+  vv = float(0.0)
+
+  for j in range(1, dim):
+    icj = map_efc2iefc_in[worldid, contact_efc_address_in[conid, j]]
+    jaj = Jaref_in[worldid, icj]
+    jvj = jv_in[worldid, icj]
+    dj = iefc_D_in[worldid, icj]
+    DJj = dj * jaj
+
+    quad += wp.vec3(0.5 * jaj * DJj, jvj * DJj, 0.5 * jvj * dj * jvj)
+
+    frictionj = friction[j - 1]
+    uj = jaj * frictionj
+    vj = jvj * frictionj
+    uu += uj * uj
+    uv += uj * vj
+    vv += vj * vj
+
+  mu2 = mu * mu
+  dm = math.safe_div(D0, mu2 * (1.0 + mu2))
+
+  # Compute N, Tsqr
+  N = u0 + alpha * v0
+  Tsqr = uu + alpha * (2.0 * uv + alpha * vv)
+
+  if Tsqr <= 0.0:
+    if N < 0.0:
+      return _eval_pt(quad, alpha)
+    return wp.vec3(0.0)
+
+  T = wp.sqrt(Tsqr)
+
+  # top zone
+  if N >= mu * T:
+    return wp.vec3(0.0)
+  # bottom zone
+  if mu * N + T <= 0.0:
+    return _eval_pt(quad, alpha)
+
+  # middle zone
+  N1 = v0
+  T1 = (uv + alpha * vv) / T
+  T2 = vv / T - (uv + alpha * vv) * T1 / (T * T)
+  nmt = N - mu * T
+  return wp.vec3(
+    0.5 * dm * nmt * nmt,
+    dm * nmt * (N1 - mu * T1),
+    dm * ((N1 - mu * T1) * (N1 - mu * T1) + nmt * (-mu * T2)),
+  )
+
+
+# TODO(team): refactor linesearch_island
+@wp.kernel
+def linesearch_island(
+  # Model:
+  opt_tolerance: wp.array[float],
+  opt_ls_tolerance: wp.array[float],
+  opt_ls_iterations: int,
+  opt_impratio_invsqrt: wp.array[float],
+  stat_meaninertia: wp.array[float],
+  # Data in:
+  nefc_in: wp.array[int],
+  nisland_in: wp.array[int],
+  contact_friction_in: wp.array[types.vec5],
+  contact_dim_in: wp.array[int],
+  contact_efc_address_in: wp.array2d[int],
+  nidof_in: wp.array[int],
+  island_nv_in: wp.array2d[int],
+  island_ne_in: wp.array2d[int],
+  island_nf_in: wp.array2d[int],
+  island_efcadr_in: wp.array2d[int],
+  island_nefc_in: wp.array2d[int],
+  map_efc2iefc_in: wp.array2d[int],
+  njmax_in: int,
+  nacon_in: wp.array[int],
+  island_idofadr_in: wp.array2d[int],
+  # In:
+  iefc_type_in: wp.array2d[int],
+  iefc_id_in: wp.array2d[int],
+  iefc_D_in: wp.array2d[float],
+  iefc_frictionloss_in: wp.array2d[float],
+  Jaref_in: wp.array2d[float],
+  jv_in: wp.array2d[float],
+  mv_in: wp.array2d[float],
+  search_in: wp.array2d[float],
+  ifrc_smooth_in: wp.array2d[float],
+  iMa_in: wp.array2d[float],
+  island_search_dot_in: wp.array2d[float],
+  island_gauss_in: wp.array2d[float],
+  island_done_in: wp.array2d[bool],
+  # Out:
+  island_alpha_out: wp.array2d[float],
+):
+  """Linesearch per island."""
+  worldid, islandid = wp.tid()
+  nisland = nisland_in[worldid]
+  if islandid >= nisland:
+    island_alpha_out[worldid, islandid] = 0.0
+    return
+  nefc = wp.min(njmax_in, nefc_in[worldid])
+  tolerance = opt_tolerance[worldid % opt_tolerance.shape[0]]
+  ls_tolerance = opt_ls_tolerance[worldid % opt_ls_tolerance.shape[0]]
+  meaninertia = stat_meaninertia[worldid % stat_meaninertia.shape[0]]
+  impratio_invsqrt = opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]]
+
+  if island_done_in[worldid, islandid]:
+    island_alpha_out[worldid, islandid] = 0.0
+    return
+
+  iefcadr = island_efcadr_in[worldid, islandid]
+  ine = island_ne_in[worldid, islandid]
+  inf = island_nf_in[worldid, islandid]
+  inv = island_nv_in[worldid, islandid]
+  idofadr = island_idofadr_in[worldid, islandid]
+
+  # Get island nefc
+  isle_nefc_end = iefcadr + island_nefc_in[worldid, islandid]
+
+  # Compute gauss quad: [gauss, s.T @ (Ma - frc_smooth), 0.5 * s.T @ mv]
+  quad_gauss_1 = float(0.0)
+  quad_gauss_2 = float(0.0)
+  for i in range(inv):
+    idof = idofadr + i
+    s = search_in[worldid, idof]
+    quad_gauss_1 += s * (iMa_in[worldid, idof] - ifrc_smooth_in[worldid, idof])
+    quad_gauss_2 += 0.5 * s * mv_in[worldid, idof]
+  quad_gauss = wp.vec3(island_gauss_in[worldid, islandid], quad_gauss_1, quad_gauss_2)
+
+  # gtol
+  snorm = wp.sqrt(island_search_dot_in[worldid, islandid])
+  scale = meaninertia * float(inv)
+  gtol = wp.max(tolerance * ls_tolerance * snorm * scale, 1e-6)
+
+  # p0: cost/grad/hessian at alpha=0
+  p0 = wp.vec3(quad_gauss[0], quad_gauss[1], 2.0 * quad_gauss[2])
+  for iefcid in range(iefcadr, isle_nefc_end):
+    if iefcid >= nefc:
+      break
+    local_iefcid = iefcid - iefcadr
+    D = iefc_D_in[worldid, iefcid]
+    ja = Jaref_in[worldid, iefcid]
+    jv_val = jv_in[worldid, iefcid]
+    if local_iefcid < ine:
+      # Equality: always active
+      jvD = jv_val * D
+      p0 += wp.vec3(0.5 * D * ja * ja, jvD * ja, jv_val * jvD)
+    elif local_iefcid < ine + inf:
+      # Friction
+      f = iefc_frictionloss_in[worldid, iefcid]
+      rf = math.safe_div(f, D)
+      p0 += _eval_frictionloss_pt(ja, f, rf, jv_val, D)
+    elif iefc_type_in[worldid, iefcid] == types.ConstraintType.CONTACT_ELLIPTIC:
+      conid = iefc_id_in[worldid, iefcid]
+      if conid < nacon_in[0]:
+        ic0 = map_efc2iefc_in[worldid, contact_efc_address_in[conid, 0]]
+        if iefcid == ic0:
+          p0 += _eval_elliptic_cost_island(
+            impratio_invsqrt,
+            contact_friction_in,
+            contact_dim_in,
+            contact_efc_address_in,
+            map_efc2iefc_in,
+            0.0,
+            conid,
+            iefc_D_in,
+            Jaref_in,
+            jv_in,
+            worldid,
+          )
+    else:
+      # Inequality
+      if ja < 0.0:
+        jvD = jv_val * D
+        p0 += wp.vec3(0.5 * D * ja * ja, jvD * ja, jv_val * jvD)
+
+  # Newton step: lo_alpha_in = -p0[1] / p0[2]
+  lo_alpha_in = -math.safe_div(p0[1], p0[2])
+
+  # Evaluate at Newton step
+  lo_in = _eval_pt(quad_gauss, lo_alpha_in)
+  for iefcid in range(iefcadr, isle_nefc_end):
+    if iefcid >= nefc:
+      break
+    local_iefcid = iefcid - iefcadr
+    D = iefc_D_in[worldid, iefcid]
+    ja = Jaref_in[worldid, iefcid]
+    jv_val = jv_in[worldid, iefcid]
+    if local_iefcid < ine:
+      lo_in += _eval_pt_direct(ja, jv_val, D, lo_alpha_in)
+    elif local_iefcid < ine + inf:
+      f = iefc_frictionloss_in[worldid, iefcid]
+      rf = math.safe_div(f, D)
+      x_a = ja + lo_alpha_in * jv_val
+      lo_in += _eval_frictionloss_pt(x_a, f, rf, jv_val, D)
+    elif iefc_type_in[worldid, iefcid] == types.ConstraintType.CONTACT_ELLIPTIC:
+      conid = iefc_id_in[worldid, iefcid]
+      if conid < nacon_in[0]:
+        ic0 = map_efc2iefc_in[worldid, contact_efc_address_in[conid, 0]]
+        if iefcid == ic0:
+          lo_in += _eval_elliptic_cost_island(
+            impratio_invsqrt,
+            contact_friction_in,
+            contact_dim_in,
+            contact_efc_address_in,
+            map_efc2iefc_in,
+            lo_alpha_in,
+            conid,
+            iefc_D_in,
+            Jaref_in,
+            jv_in,
+            worldid,
+          )
+    else:
+      x_a = ja + lo_alpha_in * jv_val
+      if x_a < 0.0:
+        lo_in += _eval_pt_direct(ja, jv_val, D, lo_alpha_in)
+
+  # Accept Newton step if derivative is small and cost improved
+  initial_converged = wp.abs(lo_in[1]) < gtol and lo_in[0] < p0[0]
+
+  if initial_converged:
+    alpha = lo_alpha_in
+  else:
+    alpha = float(0.0)
+
+    # Initialize brackets
+    lo_less = int(0)
+    if lo_in[1] < p0[1]:
+      lo_less = int(1)
+    if lo_less == 1:
+      lo = lo_in
+      lo_alpha = lo_alpha_in
+      hi = p0
+      hi_alpha = float(0.0)
+    else:
+      lo = p0
+      lo_alpha = float(0.0)
+      hi = lo_in
+      hi_alpha = lo_alpha_in
+
+    for _iter in range(opt_ls_iterations):
+      lo_next_alpha = lo_alpha - math.safe_div(lo[1], lo[2])
+      hi_next_alpha = hi_alpha - math.safe_div(hi[1], hi[2])
+      mid_alpha = 0.5 * (lo_alpha + hi_alpha)
+
+      # Evaluate at 3 candidate alphas
+      lo_next = _eval_pt(quad_gauss, lo_next_alpha)
+      hi_next = _eval_pt(quad_gauss, hi_next_alpha)
+      mid = _eval_pt(quad_gauss, mid_alpha)
+
+      for iefcid in range(iefcadr, isle_nefc_end):
+        if iefcid >= nefc:
+          break
+        local_iefcid = iefcid - iefcadr
+        D = iefc_D_in[worldid, iefcid]
+        ja = Jaref_in[worldid, iefcid]
+        jv_val = jv_in[worldid, iefcid]
+        if local_iefcid < ine:
+          r_lo, r_hi, r_mid = _eval_pt_direct_3alphas(ja, jv_val, D, lo_next_alpha, hi_next_alpha, mid_alpha)
+        elif local_iefcid < ine + inf:
+          f = iefc_frictionloss_in[worldid, iefcid]
+          rf = math.safe_div(f, D)
+          x_lo = ja + lo_next_alpha * jv_val
+          x_hi = ja + hi_next_alpha * jv_val
+          x_mid = ja + mid_alpha * jv_val
+          r_lo = _eval_frictionloss_pt(x_lo, f, rf, jv_val, D)
+          r_hi = _eval_frictionloss_pt(x_hi, f, rf, jv_val, D)
+          r_mid = _eval_frictionloss_pt(x_mid, f, rf, jv_val, D)
+        elif iefc_type_in[worldid, iefcid] == types.ConstraintType.CONTACT_ELLIPTIC:
+          conid = iefc_id_in[worldid, iefcid]
+          r_lo = wp.vec3(0.0)
+          r_hi = wp.vec3(0.0)
+          r_mid = wp.vec3(0.0)
+          if conid < nacon_in[0]:
+            ic0 = map_efc2iefc_in[worldid, contact_efc_address_in[conid, 0]]
+            if iefcid == ic0:
+              r_lo = _eval_elliptic_cost_island(
+                impratio_invsqrt,
+                contact_friction_in,
+                contact_dim_in,
+                contact_efc_address_in,
+                map_efc2iefc_in,
+                lo_next_alpha,
+                conid,
+                iefc_D_in,
+                Jaref_in,
+                jv_in,
+                worldid,
+              )
+              r_hi = _eval_elliptic_cost_island(
+                impratio_invsqrt,
+                contact_friction_in,
+                contact_dim_in,
+                contact_efc_address_in,
+                map_efc2iefc_in,
+                hi_next_alpha,
+                conid,
+                iefc_D_in,
+                Jaref_in,
+                jv_in,
+                worldid,
+              )
+              r_mid = _eval_elliptic_cost_island(
+                impratio_invsqrt,
+                contact_friction_in,
+                contact_dim_in,
+                contact_efc_address_in,
+                map_efc2iefc_in,
+                mid_alpha,
+                conid,
+                iefc_D_in,
+                Jaref_in,
+                jv_in,
+                worldid,
+              )
+        else:
+          x_lo = ja + lo_next_alpha * jv_val
+          x_hi = ja + hi_next_alpha * jv_val
+          x_mid = ja + mid_alpha * jv_val
+          r_lo = wp.vec3(0.0)
+          r_hi = wp.vec3(0.0)
+          r_mid = wp.vec3(0.0)
+          if x_lo < 0.0:
+            r_lo = _eval_pt_direct(ja, jv_val, D, lo_next_alpha)
+          if x_hi < 0.0:
+            r_hi = _eval_pt_direct(ja, jv_val, D, hi_next_alpha)
+          if x_mid < 0.0:
+            r_mid = _eval_pt_direct(ja, jv_val, D, mid_alpha)
+        lo_next += r_lo
+        hi_next += r_hi
+        mid += r_mid
+
+      # Bracket swapping
+      swap_lo = int(0)
+      if _in_bracket(lo, lo_next):
+        lo = lo_next
+        lo_alpha = lo_next_alpha
+        swap_lo = int(1)
+      if _in_bracket(lo, mid):
+        lo = mid
+        lo_alpha = mid_alpha
+        swap_lo = int(1)
+      if _in_bracket(lo, hi_next):
+        lo = hi_next
+        lo_alpha = hi_next_alpha
+        swap_lo = int(1)
+
+      swap_hi = int(0)
+      if _in_bracket(hi, hi_next):
+        hi = hi_next
+        hi_alpha = hi_next_alpha
+        swap_hi = int(1)
+      if _in_bracket(hi, mid):
+        hi = mid
+        hi_alpha = mid_alpha
+        swap_hi = int(1)
+      if _in_bracket(hi, lo_next):
+        hi = lo_next
+        hi_alpha = lo_next_alpha
+        swap_hi = int(1)
+
+      # Done check
+      ls_done = (swap_lo == 0 and swap_hi == 0) or (lo[1] < 0.0 and lo[1] > -gtol) or (hi[1] > 0.0 and hi[1] < gtol)
+
+      # Update alpha if improved
+      if lo[0] < p0[0] or hi[0] < p0[0]:
+        if lo[0] < hi[0]:
+          alpha = lo_alpha
+        else:
+          alpha = hi_alpha
+
+      if ls_done:
+        break
+
+  island_alpha_out[worldid, islandid] = alpha
+
+
+@wp.kernel
+def linesearch_qacc_ma_island(
+  # Data in:
+  nidof_in: wp.array[int],
+  # In:
+  search_in: wp.array2d[float],
+  mv_in: wp.array2d[float],
+  island_alpha_in: wp.array2d[float],
+  idof_islandid_in: wp.array2d[int],
+  island_done_in: wp.array2d[bool],
+  # Out:
+  iacc_out: wp.array2d[float],
+  iMa_out: wp.array2d[float],
+):
+  """Update iacc and iMa after linesearch."""
+  worldid, tid = wp.tid()
+
+  # Process DOFs
+  if tid < nidof_in[worldid]:
+    idof = tid
+    islandid = idof_islandid_in[worldid, idof]
+    if islandid >= 0 and not island_done_in[worldid, islandid]:
+      alpha = island_alpha_in[worldid, islandid]
+      iacc_out[worldid, idof] += alpha * search_in[worldid, idof]
+      iMa_out[worldid, idof] += alpha * mv_in[worldid, idof]
+
+
+@wp.kernel
+def linesearch_jaref_island(
+  # Data in:
+  nefc_in: wp.array[int],
+  njmax_in: int,
+  # In:
+  jv_in: wp.array2d[float],
+  island_alpha_in: wp.array2d[float],
+  iefc_islandid_in: wp.array2d[int],
+  island_done_in: wp.array2d[bool],
+  # Out:
+  Jaref_out: wp.array2d[float],
+):
+  """Update Jaref after linesearch."""
+  worldid, iefcid = wp.tid()
+
+  if iefcid >= wp.min(njmax_in, nefc_in[worldid]):
+    return
+
+  islandid = iefc_islandid_in[worldid, iefcid]
+  if islandid < 0:
+    return
+  if island_done_in[worldid, islandid]:
+    return
+
+  alpha = island_alpha_in[worldid, islandid]
+  Jaref_out[worldid, iefcid] += alpha * jv_in[worldid, iefcid]
+
+
+@wp.kernel
+def solve_prev_grad_Mgrad_island(
+  # Data in:
+  nidof_in: wp.array[int],
+  # In:
+  grad_in: wp.array2d[float],
+  Mgrad_in: wp.array2d[float],
+  idof_islandid_in: wp.array2d[int],
+  island_done_in: wp.array2d[bool],
+  # Out:
+  prev_grad_out: wp.array2d[float],
+  prev_Mgrad_out: wp.array2d[float],
+):
+  """Save prev_grad and prev_Mgrad per island DOF."""
+  worldid, idofid = wp.tid()
+
+  if idofid >= nidof_in[worldid]:
+    return
+
+  islandid = idof_islandid_in[worldid, idofid]
+  if islandid < 0:
+    return
+  if island_done_in[worldid, islandid]:
+    return
+
+  prev_grad_out[worldid, idofid] = grad_in[worldid, idofid]
+  prev_Mgrad_out[worldid, idofid] = Mgrad_in[worldid, idofid]
+
+
+@wp.kernel
+def solve_beta_island(
+  # Data in:
+  nisland_in: wp.array[int],
+  island_nv_in: wp.array2d[int],
+  # In:
+  island_idofadr_in: wp.array2d[int],
+  grad_in: wp.array2d[float],
+  Mgrad_in: wp.array2d[float],
+  prev_grad_in: wp.array2d[float],
+  prev_Mgrad_in: wp.array2d[float],
+  island_done_in: wp.array2d[bool],
+  # Out:
+  island_beta_out: wp.array2d[float],
+):
+  """Polak-Ribière beta per island."""
+  worldid, islandid = wp.tid()
+
+  if islandid >= nisland_in[worldid]:
+    return
+
+  if island_done_in[worldid, islandid]:
+    island_beta_out[worldid, islandid] = 0.0
+    return
+
+  idofadr = island_idofadr_in[worldid, islandid]
+  inv = island_nv_in[worldid, islandid]
+
+  beta_num = float(0.0)
+  beta_den = float(0.0)
+  for i in range(inv):
+    idof = idofadr + i
+    pMg = prev_Mgrad_in[worldid, idof]
+    beta_num += grad_in[worldid, idof] * (Mgrad_in[worldid, idof] - pMg)
+    beta_den += prev_grad_in[worldid, idof] * pMg
+
+  island_beta_out[worldid, islandid] = wp.max(0.0, beta_num / wp.max(types.MJ_MINVAL, beta_den))
+
+
+@wp.kernel
+def solve_search_update_island(
+  # Model:
+  opt_solver: int,
+  # Data in:
+  nidof_in: wp.array[int],
+  # In:
+  Mgrad_in: wp.array2d[float],
+  search_in: wp.array2d[float],
+  island_beta_in: wp.array2d[float],
+  idof_islandid_in: wp.array2d[int],
+  island_done_in: wp.array2d[bool],
+  # Out:
+  search_out: wp.array2d[float],
+  island_search_dot_out: wp.array2d[float],
+):
+  """Update search direction per island DOF."""
+  worldid, idofid = wp.tid()
+
+  if idofid >= nidof_in[worldid]:
+    return
+
+  islandid = idof_islandid_in[worldid, idofid]
+  if islandid < 0:
+    return
+  if island_done_in[worldid, islandid]:
+    return
+
+  s = -Mgrad_in[worldid, idofid]
+  if opt_solver == types.SolverType.CG:
+    s += island_beta_in[worldid, islandid] * search_in[worldid, idofid]
+
+  search_out[worldid, idofid] = s
+  wp.atomic_add(island_search_dot_out, worldid, islandid, s * s)
+
+
+@wp.kernel
+def solve_init_nsolving_island(
+  nisland_in: wp.array[int],
+  nsolving_out: wp.array[int],
+):
+  """Initialize active island count nsolving on device without CPU stalls."""
+  worldid = wp.tid()
+  wp.atomic_add(nsolving_out, 0, nisland_in[worldid])
+
+
+@wp.kernel
+def solve_done_island(
+  # Model:
+  opt_tolerance: wp.array[float],
+  opt_iterations: int,
+  stat_meaninertia: wp.array[float],
+  # Data in:
+  nisland_in: wp.array[int],
+  island_nv_in: wp.array2d[int],
+  # In:
+  island_grad_dot_in: wp.array2d[float],
+  island_cost_in: wp.array2d[float],
+  island_prev_cost_in: wp.array2d[float],
+  island_done_in: wp.array2d[bool],
+  # Data out:
+  solver_niter_out: wp.array[int],
+  # Out:
+  island_done_out: wp.array2d[bool],
+  island_solver_niter_out: wp.array2d[int],
+  nsolving_out: wp.array[int],
+):
+  """Check convergence per island."""
+  worldid, islandid = wp.tid()
+
+  if islandid >= nisland_in[worldid]:
+    return
+
+  tolerance = opt_tolerance[worldid % opt_tolerance.shape[0]]
+  meaninertia = stat_meaninertia[worldid % stat_meaninertia.shape[0]]
+
+  if island_done_in[worldid, islandid]:
+    niter = island_solver_niter_out[worldid, islandid]
+    wp.atomic_max(solver_niter_out, worldid, niter)
+    return
+
+  island_solver_niter_out[worldid, islandid] += 1
+  niter = island_solver_niter_out[worldid, islandid]
+  wp.atomic_max(solver_niter_out, worldid, niter)
+
+  inv = island_nv_in[worldid, islandid]
+  improvement = _rescale(inv, meaninertia, island_prev_cost_in[worldid, islandid] - island_cost_in[worldid, islandid])
+  gradient = _rescale(inv, meaninertia, wp.sqrt(island_grad_dot_in[worldid, islandid]))
+  done = (improvement < tolerance) or (gradient < tolerance)
+  if done or niter >= opt_iterations:
+    island_done_out[worldid, islandid] = True
+    wp.atomic_sub(nsolving_out, 0, 1)
+
+
+@wp.kernel
+def update_gradient_JTDAJ_island(
+  # Model:
+  is_sparse: bool,
+  # Data in:
+  nefc_in: wp.array[int],
+  njmax_in: int,
+  island_idofadr_in: wp.array2d[int],
+  island_nv_in: wp.array2d[int],
+  # In:
+  iefc_J_rownnz_in: wp.array2d[int],
+  iefc_J_rowadr_in: wp.array2d[int],
+  iefc_J_colind_in: wp.array3d[int],
+  iefc_J_in: wp.array3d[float],
+  iefc_D_in: wp.array2d[float],
+  iefc_state_in: wp.array2d[int],
+  iefc_islandid_in: wp.array2d[int],
+  island_done_in: wp.array2d[bool],
+  # Out:
+  ih_out: wp.array3d[float],
+):
+  """Build island Hessian: ih += Jᵀ·D·J for active constraints."""
+  worldid, iefcid = wp.tid()
+
+  if iefcid >= wp.min(njmax_in, nefc_in[worldid]):
+    return
+
+  islandid = iefc_islandid_in[worldid, iefcid]
+  if islandid < 0:
+    return
+  if island_done_in[worldid, islandid]:
+    return
+
+  state = iefc_state_in[worldid, iefcid]
+  if state != types.ConstraintState.QUADRATIC.value:
+    return
+  D = iefc_D_in[worldid, iefcid]
+
+  idofadr = island_idofadr_in[worldid, islandid]
+  inv = island_nv_in[worldid, islandid]
+  if is_sparse:
+    rownnz = iefc_J_rownnz_in[worldid, iefcid]
+    rowadr = iefc_J_rowadr_in[worldid, iefcid]
+    for k1 in range(rownnz):
+      adr1 = rowadr + k1
+      Ji = iefc_J_in[worldid, 0, adr1]
+      i = iefc_J_colind_in[worldid, 0, adr1]
+
+      for k2 in range(k1 + 1):
+        adr2 = rowadr + k2
+        Jj = iefc_J_in[worldid, 0, adr2]
+        j = iefc_J_colind_in[worldid, 0, adr2]
+
+        h = Ji * Jj * D
+        wp.atomic_add(ih_out[worldid, i], j, h)
+        if i != j:
+          wp.atomic_add(ih_out[worldid, j], i, h)
+  else:
+    for ii in range(inv):
+      i = idofadr + ii
+      Ji = iefc_J_in[worldid, iefcid, i]
+      if Ji == 0.0:
+        continue
+      for jj in range(ii + 1):
+        j = idofadr + jj
+        Jj = iefc_J_in[worldid, iefcid, j]
+        if Jj == 0.0:
+          continue
+        h = Ji * Jj * D
+        wp.atomic_add(ih_out[worldid, i], j, h)
+        if i != j:
+          wp.atomic_add(ih_out[worldid, j], i, h)
+
+
+@wp.kernel
+def update_gradient_set_h_M_sparse_island(
+  # Model:
+  M_fullm_i: wp.array[int],
+  M_fullm_j: wp.array[int],
+  M_elemid: wp.array2d[int],
+  # Data in:
+  nidof_in: wp.array[int],
+  M_in: wp.array3d[float],
+  dof_island_in: wp.array2d[int],
+  map_dof2idof_in: wp.array2d[int],
+  # In:
+  island_done_in: wp.array2d[bool],
+  # Out:
+  ih_out: wp.array3d[float],
+):
+  """Add sparse mass matrix to island Hessian using global-to-island DOF mapping."""
+  worldid, elementid = wp.tid()
+
+  i_global = M_fullm_i[elementid]
+  j_global = M_fullm_j[elementid]
+
+  madr = M_elemid[i_global, j_global]
+  if madr < 0:
+    return
+
+  # Check both DOFs belong to an island
+  island_i = dof_island_in[worldid, i_global]
+  if island_i < 0:
+    return
+  if island_done_in[worldid, island_i]:
+    return
+
+  island_j = dof_island_in[worldid, j_global]
+  if island_j < 0:
+    return
+
+  # Both DOFs must be in the same island
+  if island_i != island_j:
+    return
+
+  idof_i = map_dof2idof_in[worldid, i_global]
+  idof_j = map_dof2idof_in[worldid, j_global]
+
+  val = M_in[worldid, 0, madr]
+  ih_out[worldid, idof_i, idof_j] += val
+  if idof_i != idof_j:
+    ih_out[worldid, idof_j, idof_i] += val
+
+
+@wp.kernel
+def update_gradient_set_h_M_dense_island(
+  # Model:
+  nv: int,
+  # Data in:
+  nidof_in: wp.array[int],
+  M_in: wp.array3d[float],
+  map_idof2dof_in: wp.array2d[int],
+  # In:
+  idof_islandid_in: wp.array2d[int],
+  island_done_in: wp.array2d[bool],
+  # Out:
+  ih_out: wp.array3d[float],
+):
+  """Add dense mass matrix to island Hessian."""
+  worldid, idofid = wp.tid()
+
+  if idofid >= nidof_in[worldid]:
+    return
+
+  islandid = idof_islandid_in[worldid, idofid]
+  if islandid < 0:
+    return
+  if island_done_in[worldid, islandid]:
+    return
+
+  dof_i = map_idof2dof_in[worldid, idofid]
+
+  # Copy row from M to ih, mapping columns
+  nid = nidof_in[worldid]
+  for jdof in range(nid):
+    dof_j = map_idof2dof_in[worldid, jdof]
+    ih_out[worldid, idofid, jdof] += M_in[worldid, dof_i, dof_j]
+
+
+@wp.kernel
+def update_gradient_JTCJ_island(
+  # Model:
+  opt_impratio_invsqrt: wp.array[float],
+  is_sparse: bool,
+  # Data in:
+  nacon_in: wp.array[int],
+  contact_friction_in: wp.array[types.vec5],
+  contact_dim_in: wp.array[int],
+  contact_efc_address_in: wp.array2d[int],
+  contact_worldid_in: wp.array[int],
+  island_idofadr_in: wp.array2d[int],
+  naconmax_in: int,
+  nidof_in: wp.array[int],
+  map_efc2iefc_in: wp.array2d[int],
+  island_nv_in: wp.array2d[int],
+  # In:
+  iefc_J_rownnz_in: wp.array2d[int],
+  iefc_J_rowadr_in: wp.array2d[int],
+  iefc_J_colind_in: wp.array3d[int],
+  iefc_J_in: wp.array3d[float],
+  iefc_D_in: wp.array2d[float],
+  iefc_state_in: wp.array2d[int],
+  Jaref_in: wp.array2d[float],
+  iefc_islandid_in: wp.array2d[int],
+  island_done_in: wp.array2d[bool],
+  # Out:
+  ih_out: wp.array3d[float],
+):
+  """Add elliptic cone Hessian correction: Jᵀ·C·J for contacts in CONE state."""
+  conid = wp.tid()
+
+  if conid >= wp.min(naconmax_in, nacon_in[0]):
+    return
+
+  worldid = contact_worldid_in[conid]
+  condim = contact_dim_in[conid]
+
+  if condim == 1:
+    return
+
+  efcid0_global = contact_efc_address_in[conid, 0]
+  if efcid0_global < 0:
+    return
+
+  ic0 = map_efc2iefc_in[worldid, efcid0_global]
+  if iefc_state_in[worldid, ic0] != types.ConstraintState.CONE.value:
+    return
+
+  islandid = iefc_islandid_in[worldid, ic0]
+  if islandid < 0:
+    return
+  if island_done_in[worldid, islandid]:
+    return
+
+  inv = island_nv_in[worldid, islandid]
+  idofadr = island_idofadr_in[worldid, islandid]
+
+  fri = contact_friction_in[conid]
+  mu = fri[0] * opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]]
+  mu2 = mu * mu
+  dm = math.safe_div(iefc_D_in[worldid, ic0], mu2 * (1.0 + mu2))
+
+  if dm == 0.0:
+    return
+
+  # Compute n and u vector
+  n = Jaref_in[worldid, ic0] * mu
+  u = types.vec6(n, 0.0, 0.0, 0.0, 0.0, 0.0)
+
+  tt = float(0.0)
+  for j in range(1, condim):
+    efcidj_global = contact_efc_address_in[conid, j]
+    if efcidj_global < 0:
+      return
+    icj = map_efc2iefc_in[worldid, efcidj_global]
+    uj = Jaref_in[worldid, icj] * fri[j - 1]
+    tt += uj * uj
+    u[j] = uj
+
+  if tt <= 0.0:
+    t = 0.0
+  else:
+    t = wp.sqrt(tt)
+  t = wp.max(t, types.MJ_MINVAL)
+  ttt = wp.max(t * t * t, types.MJ_MINVAL)
+
+  # Accumulate cone correction into ih
+  for dim1id in range(condim):
+    if dim1id == 0:
+      ic1 = ic0
+    else:
+      efcid1_global = contact_efc_address_in[conid, dim1id]
+      if efcid1_global < 0:
+        return
+      ic1 = map_efc2iefc_in[worldid, efcid1_global]
+
+    ui = u[dim1id]
+
+    for dim2id in range(dim1id + 1):
+      if dim2id == 0:
+        ic2 = ic0
+      else:
+        efcid2_global = contact_efc_address_in[conid, dim2id]
+        if efcid2_global < 0:
+          return
+        ic2 = map_efc2iefc_in[worldid, efcid2_global]
+
+      uj = u[dim2id]
+
+      # Cone correction matrix
+      if dim1id == 0 and dim2id == 0:
+        hcone = 1.0
+      elif dim1id == 0:
+        hcone = -math.safe_div(mu, t) * uj
+      elif dim2id == 0:
+        hcone = -math.safe_div(mu, t) * ui
+      else:
+        hcone = mu * math.safe_div(n, ttt) * ui * uj
+        if dim1id == dim2id:
+          hcone += mu2 - mu * math.safe_div(n, t)
+
+      # Scale by dm * friction
+      if dim1id == 0:
+        fri1 = mu
+      else:
+        fri1 = fri[dim1id - 1]
+      if dim2id == 0:
+        fri2 = mu
+      else:
+        fri2 = fri[dim2id - 1]
+
+      hcone *= dm * fri1 * fri2
+
+      if hcone == 0.0:
+        continue
+
+      # Accumulate J1^T * hcone * J2 into ih (lower triangle)
+      if is_sparse:
+        if dim1id == dim2id:
+          rownnz = iefc_J_rownnz_in[worldid, ic1]
+          rowadr = iefc_J_rowadr_in[worldid, ic1]
+          for k1 in range(rownnz):
+            adr1 = rowadr + k1
+            J1 = iefc_J_in[worldid, 0, adr1]
+            i = iefc_J_colind_in[worldid, 0, adr1]
+
+            for k2 in range(k1 + 1):
+              adr2 = rowadr + k2
+              J2 = iefc_J_in[worldid, 0, adr2]
+              j = iefc_J_colind_in[worldid, 0, adr2]
+
+              val = hcone * J1 * J2
+              wp.atomic_add(ih_out[worldid, i], j, val)
+              if i != j:
+                wp.atomic_add(ih_out[worldid, j], i, val)
+        else:
+          rownnz1 = iefc_J_rownnz_in[worldid, ic1]
+          rowadr1 = iefc_J_rowadr_in[worldid, ic1]
+          rownnz2 = iefc_J_rownnz_in[worldid, ic2]
+          rowadr2 = iefc_J_rowadr_in[worldid, ic2]
+
+          for k1 in range(rownnz1):
+            adr1 = rowadr1 + k1
+            J1 = iefc_J_in[worldid, 0, adr1]
+            i = iefc_J_colind_in[worldid, 0, adr1]
+
+            for k2 in range(rownnz2):
+              adr2 = rowadr2 + k2
+              J2 = iefc_J_in[worldid, 0, adr2]
+              j = iefc_J_colind_in[worldid, 0, adr2]
+
+              val = hcone * J1 * J2
+              if i == j:
+                wp.atomic_add(ih_out[worldid, i], j, val * 2.0)
+              else:
+                wp.atomic_add(ih_out[worldid, i], j, val)
+                wp.atomic_add(ih_out[worldid, j], i, val)
+      else:
+        for i in range(inv):
+          J1i = iefc_J_in[worldid, ic1, idofadr + i]
+          if J1i == 0.0:
+            continue
+          for jj in range(i + 1):
+            J2j = iefc_J_in[worldid, ic2, idofadr + jj]
+            if J2j == 0.0:
+              continue
+            val = hcone * J1i * J2j
+            wp.atomic_add(ih_out[worldid, idofadr + i], idofadr + jj, val)
+            if i != jj:
+              wp.atomic_add(ih_out[worldid, idofadr + jj], idofadr + i, val)
+
+            if dim1id != dim2id:
+              J1j = iefc_J_in[worldid, ic1, idofadr + jj]
+              J2i = iefc_J_in[worldid, ic2, idofadr + i]
+              if J1j != 0.0 and J2i != 0.0:
+                val2 = hcone * J1j * J2i
+                wp.atomic_add(ih_out[worldid, idofadr + i], idofadr + jj, val2)
+                if i != jj:
+                  wp.atomic_add(ih_out[worldid, idofadr + jj], idofadr + i, val2)
+
+
+@wp.kernel
+def _cholesky_factorize_solve_island(
+  # Data in:
+  nisland_in: wp.array[int],
+  island_idofadr_in: wp.array2d[int],
+  island_nv_in: wp.array2d[int],
+  # In:
+  grad_in: wp.array2d[float],
+  ih_in: wp.array3d[float],
+  island_done_in: wp.array2d[bool],
+  # Out:
+  Mgrad_out: wp.array2d[float],
+):
+  """Per-island Cholesky factorize and solve: Mgrad = H⁻¹ @ grad.
+
+  One thread per (world, island). Performs dense in-place Cholesky factorization
+  on the island's inv x inv subblock of ih, then forward/backward substitution.
+  """
+  worldid, islandid = wp.tid()
+
+  if islandid >= nisland_in[worldid]:
+    return
+  if island_done_in[worldid, islandid]:
+    return
+
+  inv = island_nv_in[worldid, islandid]
+
+  if inv == 0:
+    return
+
+  adr = island_idofadr_in[worldid, islandid]
+  # Cholesky factorization in-place: L such that H = L @ L^T
+  for i in range(inv):
+    for j in range(i + 1):
+      s = ih_in[worldid, adr + i, adr + j]
+      for k in range(j):
+        s -= ih_in[worldid, adr + i, adr + k] * ih_in[worldid, adr + j, adr + k]
+      if i == j:
+        if s <= 1e-6:
+          s = 1e-6
+        ih_in[worldid, adr + i, adr + j] = wp.sqrt(s)
+      else:
+        div = ih_in[worldid, adr + j, adr + j]
+        ih_in[worldid, adr + i, adr + j] = s / wp.max(1e-6, div)
+
+  # Forward substitution: L @ y = grad  =>  y
+  for i in range(inv):
+    s = grad_in[worldid, adr + i]
+    for k in range(i):
+      s -= ih_in[worldid, adr + i, adr + k] * Mgrad_out[worldid, adr + k]
+    Mgrad_out[worldid, adr + i] = s / wp.max(1e-6, ih_in[worldid, adr + i, adr + i])
+
+  # Backward substitution: L^T @ x = y  =>  x = Mgrad
+  for i_rev in range(inv):
+    i = inv - 1 - i_rev
+    s = Mgrad_out[worldid, adr + i]
+    for k in range(i + 1, inv):
+      s -= ih_in[worldid, adr + k, adr + i] * Mgrad_out[worldid, adr + k]
+    Mgrad_out[worldid, adr + i] = s / wp.max(types.MJ_MINVAL, ih_in[worldid, adr + i, adr + i])
+
+
+def init_context_island(m: types.Model, d: types.Data, ctx: IslandSolverContext):
+  """Initialize island solver context."""
+  # Init per-island scalars
+  d.solver_niter.zero_()
+  wp.launch(
+    solve_init_efc_island,
+    dim=(d.nworld, m.ntree),
+    inputs=[d.nisland],
+    outputs=[ctx.cost, ctx.search_dot, ctx.done, ctx.solver_niter],
+  )
+
+  # Jaref = iefc_J @ iacc - iefc_aref
+  wp.launch(
+    solve_init_jaref_island,
+    dim=(d.nworld, d.njmax),
+    inputs=[
+      m.is_sparse,
+      d.nefc,
+      d.island_nv,
+      d.njmax,
+      d.island_dofadr,
+      d.efc.iJ_rownnz,
+      d.efc.iJ_rowadr,
+      d.efc.iJ_colind,
+      d.efc.iJ,
+      d.iqacc,
+      d.efc.iaref,
+      d.efc_islandid,
+      ctx.done,
+    ],
+    outputs=[ctx.Jaref],
+  )
+
+  # iMa = M @ iacc (all islands in parallel)
+  support.mul_m_island(
+    m,
+    d,
+    ctx.Ma,
+    d.iqacc,
+    d.nidof,
+    d.map_idof2dof,
+    d.map_dof2idof,
+    d.dof_islandid,
+  )
+
+  # Update constraint
+  _update_constraint_island(m, d, ctx)
+
+  # Update gradient
+  if m.opt.solver == types.SolverType.NEWTON:
+    _update_gradient_incremental_island(m, d, ctx)
+  else:
+    _update_gradient_island(m, d, ctx)
+
+
+@event_scope
+def _update_constraint_island(m: types.Model, d: types.Data, ctx: IslandSolverContext):
+  """Update constraint arrays for island solver."""
+  # Save prev cost, zero cost/gauss
+  wp.launch(
+    update_constraint_init_cost_island,
+    dim=(d.nworld, m.ntree),
+    inputs=[d.nisland, ctx.cost, ctx.done],
+    outputs=[ctx.gauss, ctx.cost, ctx.prev_cost],
+  )
+
+  # Compute force, state, cost per EFC
+  wp.launch(
+    update_constraint_efc_island,
+    dim=(d.nworld, d.njmax),
+    inputs=[
+      m.opt.impratio_invsqrt,
+      d.nefc,
+      d.contact.friction,
+      d.contact.dim,
+      d.contact.efc_address,
+      d.island_nefc,
+      d.island_ne,
+      d.island_nf,
+      d.island_efcadr,
+      d.map_efc2iefc,
+      d.njmax,
+      d.nacon,
+      d.efc.itype,
+      d.efc.iid,
+      d.efc.iD,
+      d.efc.ifrictionloss,
+      ctx.Jaref,
+      d.efc_islandid,
+      ctx.done,
+    ],
+    outputs=[d.efc.iforce, d.efc.istate, ctx.cost],
+  )
+
+  # qfrc_constraint = J^T @ force
+  if m.is_sparse:
+    d.iqfrc_constraint.zero_()
+    wp.launch(
+      update_constraint_init_qfrc_constraint_sparse_island,
+      dim=(d.nworld, d.njmax),
+      inputs=[
+        d.nefc,
+        d.njmax,
+        d.efc.iJ_rownnz,
+        d.efc.iJ_rowadr,
+        d.efc.iJ_colind,
+        d.efc.iJ,
+        d.efc.iforce,
+        d.efc_islandid,
+        ctx.done,
+      ],
+      outputs=[d.iqfrc_constraint],
+    )
+  else:
+    wp.launch(
+      update_constraint_init_qfrc_constraint_dense_island,
+      dim=(d.nworld, m.nv),
+      inputs=[
+        d.nefc,
+        d.nidof,
+        d.island_nefc,
+        d.island_efcadr,
+        d.njmax,
+        d.efc.iJ,
+        d.efc.iforce,
+        d.dof_islandid,
+        ctx.done,
+      ],
+      outputs=[d.iqfrc_constraint],
+    )
+
+  # Gauss cost
+  wp.launch(
+    update_constraint_gauss_cost_island,
+    dim=(d.nworld, m.nv),
+    inputs=[
+      d.nidof,
+      d.iqacc,
+      d.iqfrc_smooth,
+      d.iqacc_smooth,
+      ctx.Ma,
+      d.dof_islandid,
+      ctx.done,
+    ],
+    outputs=[ctx.gauss, ctx.cost],
+  )
+
+
+@event_scope
+def _update_gradient_island(m: types.Model, d: types.Data, ctx: IslandSolverContext):
+  """Update gradient for island solver."""
+  # Zero grad_dot per island
+  ctx.grad_dot.zero_()
+
+  # grad = Ma - frc_smooth - frc_constraint, accumulate grad_dot
+  wp.launch(
+    update_gradient_grad_island,
+    dim=(d.nworld, m.nv),
+    inputs=[
+      d.nidof,
+      d.iqfrc_smooth,
+      d.iqfrc_constraint,
+      ctx.Ma,
+      d.dof_islandid,
+      ctx.done,
+    ],
+    outputs=[ctx.grad, ctx.grad_dot],
+  )
+
+  # CG preconditioner: Mgrad = M^{-1} @ grad (direct solve)
+  support.solve_m_island(
+    m,
+    d,
+    ctx.Mgrad,
+    ctx.grad,
+    d.nidof,
+    d.map_idof2dof,
+  )
+
+
+@event_scope
+def _update_gradient_incremental_island(m: types.Model, d: types.Data, ctx: IslandSolverContext):
+  """Full Newton gradient update for islands: build H, factorize, solve."""
+  # Zero grad_dot per island
+  ctx.grad_dot.zero_()
+
+  # grad = Ma - frc_smooth - frc_constraint, accumulate grad_dot
+  wp.launch(
+    update_gradient_grad_island,
+    dim=(d.nworld, m.nv),
+    inputs=[
+      d.nidof,
+      d.iqfrc_smooth,
+      d.iqfrc_constraint,
+      ctx.Ma,
+      d.dof_islandid,
+      ctx.done,
+    ],
+    outputs=[ctx.grad, ctx.grad_dot],
+  )
+
+  # Build H = qM + Jᵀ·D·J
+  ctx.h.zero_()
+
+  # JTDAJ
+  wp.launch(
+    update_gradient_JTDAJ_island,
+    dim=(d.nworld, d.njmax),
+    inputs=[
+      m.is_sparse,
+      d.nefc,
+      d.njmax,
+      d.island_dofadr,
+      d.island_nv,
+      d.efc.iJ_rownnz,
+      d.efc.iJ_rowadr,
+      d.efc.iJ_colind,
+      d.efc.iJ,
+      d.efc.iD,
+      d.efc.istate,
+      d.efc_islandid,
+      ctx.done,
+    ],
+    outputs=[ctx.h],
+  )
+
+  # Add mass matrix
+  if m.is_sparse:
+    wp.launch(
+      update_gradient_set_h_M_sparse_island,
+      dim=(d.nworld, m.M_fullm_i.shape[0]),
+      inputs=[
+        m.M_fullm_i,
+        m.M_fullm_j,
+        m.M_elemid,
+        d.nidof,
+        d.M,
+        d.dof_island,
+        d.map_dof2idof,
+        ctx.done,
+      ],
+      outputs=[ctx.h],
+    )
+  else:
+    wp.launch(
+      update_gradient_set_h_M_dense_island,
+      dim=(d.nworld, m.nv),
+      inputs=[
+        m.nv,
+        d.nidof,
+        d.M,
+        d.map_idof2dof,
+        d.dof_islandid,
+        ctx.done,
+      ],
+      outputs=[ctx.h],
+    )
+
+  # Elliptic cone correction: JTCJ
+  if m.opt.cone == types.ConeType.ELLIPTIC and d.naconmax > 0:
+    wp.launch(
+      update_gradient_JTCJ_island,
+      dim=d.naconmax,
+      inputs=[
+        m.opt.impratio_invsqrt,
+        m.is_sparse,
+        d.nacon,
+        d.contact.friction,
+        d.contact.dim,
+        d.contact.efc_address,
+        d.contact.worldid,
+        d.island_dofadr,
+        d.naconmax,
+        d.nidof,
+        d.map_efc2iefc,
+        d.island_nv,
+        d.efc.iJ_rownnz,
+        d.efc.iJ_rowadr,
+        d.efc.iJ_colind,
+        d.efc.iJ,
+        d.efc.iD,
+        d.efc.istate,
+        ctx.Jaref,
+        d.efc_islandid,
+        ctx.done,
+      ],
+      outputs=[ctx.h],
+    )
+
+  # Cholesky factorize and solve: Mgrad = H⁻¹ @ grad
+  wp.launch(
+    _cholesky_factorize_solve_island,
+    dim=(d.nworld, m.ntree),
+    inputs=[
+      d.nisland,
+      d.island_dofadr,
+      d.island_nv,
+      ctx.grad,
+      ctx.h,
+      ctx.done,
+    ],
+    outputs=[ctx.Mgrad],
+  )
+
+
+@event_scope
+def _linesearch_island(m: types.Model, d: types.Data, ctx: IslandSolverContext):
+  """Linesearch for island solver."""
+  # mv = M @ search (all islands)
+  support.mul_m_island(
+    m,
+    d,
+    ctx.mv,
+    ctx.search,
+    d.nidof,
+    d.map_idof2dof,
+    d.map_dof2idof,
+    d.dof_islandid,
+    island_done=ctx.done,
+  )
+
+  # jv = J @ search
+  wp.launch(
+    linesearch_jv_island,
+    dim=(d.nworld, d.njmax),
+    inputs=[
+      m.is_sparse,
+      d.nefc,
+      d.nidof,
+      d.island_nv,
+      d.njmax,
+      d.island_dofadr,
+      d.efc.iJ_rownnz,
+      d.efc.iJ_rowadr,
+      d.efc.iJ_colind,
+      d.efc.iJ,
+      ctx.search,
+      d.efc_islandid,
+      ctx.done,
+    ],
+    outputs=[ctx.jv],
+  )
+
+  # linesearch
+  wp.launch(
+    linesearch_island,
+    dim=(d.nworld, m.ntree),
+    inputs=[
+      m.opt.tolerance,
+      m.opt.ls_tolerance,
+      m.opt.ls_iterations,
+      m.opt.impratio_invsqrt,
+      m.stat.meaninertia,
+      d.nefc,
+      d.nisland,
+      d.contact.friction,
+      d.contact.dim,
+      d.contact.efc_address,
+      d.nidof,
+      d.island_nv,
+      d.island_ne,
+      d.island_nf,
+      d.island_efcadr,
+      d.island_nefc,
+      d.map_efc2iefc,
+      d.njmax,
+      d.nacon,
+      d.island_dofadr,
+      d.efc.itype,
+      d.efc.iid,
+      d.efc.iD,
+      d.efc.ifrictionloss,
+      ctx.Jaref,
+      ctx.jv,
+      ctx.mv,
+      ctx.search,
+      d.iqfrc_smooth,
+      ctx.Ma,
+      ctx.search_dot,
+      ctx.gauss,
+      ctx.done,
+    ],
+    outputs=[ctx.alpha],
+  )
+
+  # Update iacc, iMa
+  wp.launch(
+    linesearch_qacc_ma_island,
+    dim=(d.nworld, m.nv),
+    inputs=[
+      d.nidof,
+      ctx.search,
+      ctx.mv,
+      ctx.alpha,
+      d.dof_islandid,
+      ctx.done,
+    ],
+    outputs=[d.iqacc, ctx.Ma],
+  )
+
+  # Update Jaref
+  wp.launch(
+    linesearch_jaref_island,
+    dim=(d.nworld, d.njmax),
+    inputs=[
+      d.nefc,
+      d.njmax,
+      ctx.jv,
+      ctx.alpha,
+      d.efc_islandid,
+      ctx.done,
+    ],
+    outputs=[ctx.Jaref],
+  )
+
+
+@event_scope
+def _solver_iteration_island(
+  m: types.Model,
+  d: types.Data,
+  ctx: IslandSolverContext,
+  nsolving: wp.array[int],
+):
+  """One iteration of island solver for all islands in parallel."""
+  _linesearch_island(m, d, ctx)
+
+  is_newton = m.opt.solver == types.SolverType.NEWTON
+  is_cg = not is_newton
+
+  # Save prev_grad, prev_Mgrad for CG
+  if is_cg:
+    wp.launch(
+      solve_prev_grad_Mgrad_island,
+      dim=(d.nworld, m.nv),
+      inputs=[d.nidof, ctx.grad, ctx.Mgrad, d.dof_islandid, ctx.done],
+      outputs=[ctx.prev_grad, ctx.prev_Mgrad],
+    )
+
+  # Update constraint
+  _update_constraint_island(m, d, ctx)
+
+  # Update gradient
+  if is_newton:
+    _update_gradient_incremental_island(m, d, ctx)
+  else:
+    _update_gradient_island(m, d, ctx)
+
+  # Polak-Ribière beta (CG only)
+  if is_cg:
+    wp.launch(
+      solve_beta_island,
+      dim=(d.nworld, m.ntree),
+      inputs=[
+        d.nisland,
+        d.island_nv,
+        d.island_dofadr,
+        ctx.grad,
+        ctx.Mgrad,
+        ctx.prev_grad,
+        ctx.prev_Mgrad,
+        ctx.done,
+      ],
+      outputs=[ctx.beta],
+    )
+
+  # Zero search_dot
+  ctx.search_dot.zero_()
+
+  # Search update
+  wp.launch(
+    solve_search_update_island,
+    dim=(d.nworld, m.nv),
+    inputs=[
+      m.opt.solver,
+      d.nidof,
+      ctx.Mgrad,
+      ctx.search,
+      ctx.beta,
+      d.dof_islandid,
+      ctx.done,
+    ],
+    outputs=[ctx.search, ctx.search_dot],
+  )
+
+  # Convergence check
+  d.solver_niter.zero_()
+  wp.launch(
+    solve_done_island,
+    dim=(d.nworld, m.ntree),
+    inputs=[
+      m.opt.tolerance,
+      m.opt.iterations,
+      m.stat.meaninertia,
+      d.nisland,
+      d.island_nv,
+      ctx.grad_dot,
+      ctx.cost,
+      ctx.prev_cost,
+      ctx.done,
+    ],
+    outputs=[d.solver_niter, ctx.done, ctx.solver_niter, nsolving],
+  )
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py
index b42e9808..61e5b219 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py
@@ -69,11 +69,11 @@ def mul_m_sparse(check_skip: bool):
   @wp.kernel(module="unique")
   def _mul_m_sparse(
     # Model:
-    qM_mulm_rowadr: wp.array[int],
-    qM_mulm_col: wp.array[int],
-    qM_mulm_madr: wp.array[int],
+    M_mulm_rowadr: wp.array[int],
+    M_mulm_col: wp.array[int],
+    M_mulm_madr: wp.array[int],
     # Data in:
-    qM_in: wp.array3d[float],
+    M_in: wp.array3d[float],
     # In:
     vec: wp.array2d[float],
     skip: wp.array[bool],
@@ -89,12 +89,12 @@ def mul_m_sparse(check_skip: bool):
 
     # Gather all contributions (diagonal + off-diagonal)
     acc = float(0.0)
-    start = qM_mulm_rowadr[dofid]
-    end = qM_mulm_rowadr[dofid + 1]
+    start = M_mulm_rowadr[dofid]
+    end = M_mulm_rowadr[dofid + 1]
     for k in range(start, end):
-      col = qM_mulm_col[k]
-      madr = qM_mulm_madr[k]
-      acc += qM_in[worldid, 0, madr] * vec[worldid, col]
+      col = M_mulm_col[k]
+      madr = M_mulm_madr[k]
+      acc += M_in[worldid, 0, madr] * vec[worldid, col]
 
     res[worldid, dofid] = acc
 
@@ -108,7 +108,7 @@ def mul_m_dense(nv: int, check_skip: bool):
   @wp.kernel(module="unique")
   def _mul_m_dense(
     # Data in:
-    qM_in: wp.array3d[float],
+    M_in: wp.array3d[float],
     # In:
     vec: wp.array2d[float],
     skip: wp.array[bool],
@@ -123,7 +123,7 @@ def mul_m_dense(nv: int, check_skip: bool):
 
     acc = float(0.0)
     for j in range(wp.static(nv)):
-      acc += qM_in[worldid, i, j] * vec[worldid, j]
+      acc += M_in[worldid, i, j] * vec[worldid, j]
     res[worldid, i] = acc
 
   return _mul_m_dense
@@ -143,8 +143,8 @@ def mul_m(
   Args:
     m: The model containing kinematic and dynamic information (device).
     d: The data object containing the current state and output arrays (device).
-    res: Result: qM @ vec.
-    vec: Input vector to multiply by qM.
+    res: Result: M @ vec.
+    vec: Input vector to multiply by M.
     skip: Per-world bitmask to skip computing output.
     M: Input matrix: M @ vec.
   """
@@ -152,13 +152,13 @@ def mul_m(
   skip = skip or wp.empty(0, dtype=bool)
 
   if M is None:
-    M = d.qM
+    M = d.M
 
   if m.is_sparse:
     wp.launch(
       mul_m_sparse(check_skip),
       dim=(d.nworld, m.nv),
-      inputs=[m.qM_mulm_rowadr, m.qM_mulm_col, m.qM_mulm_madr, M, vec, skip],
+      inputs=[m.M_mulm_rowadr, m.M_mulm_col, m.M_mulm_madr, M, vec, skip],
       outputs=[res],
     )
 
@@ -171,6 +171,340 @@ def mul_m(
     )
 
 
+@wp.kernel
+def _mul_m_island_sparse(
+  # Model:
+  M_mulm_rowadr: wp.array[int],
+  M_mulm_col: wp.array[int],
+  M_mulm_madr: wp.array[int],
+  # Data in:
+  nidof_in: wp.array[int],
+  M_in: wp.array3d[float],
+  map_dof2idof_in: wp.array2d[int],
+  map_idof2dof_in: wp.array2d[int],
+  # In:
+  idof_islandid_in: wp.array2d[int],
+  vec: wp.array2d[float],
+  island_done_in: wp.array2d[bool],
+  check_skip: int,
+  # Out:
+  res: wp.array2d[float],
+):
+  """Sparse island mul_m for ALL islands in parallel."""
+  worldid, idofid = wp.tid()
+
+  nidof = nidof_in[worldid]
+  if idofid >= nidof:
+    return
+
+  islandid = idof_islandid_in[worldid, idofid]
+  if islandid < 0:
+    return
+
+  if check_skip:
+    if island_done_in[worldid, islandid]:
+      return
+
+  dof = map_idof2dof_in[worldid, idofid]
+
+  acc = float(0.0)
+  start = M_mulm_rowadr[dof]
+  end = M_mulm_rowadr[dof + 1]
+  for k in range(start, end):
+    col = M_mulm_col[k]
+    madr = M_mulm_madr[k]
+    col_idof = map_dof2idof_in[worldid, col]
+    # skip unconstrained DOFs
+    if col_idof < nidof:
+      acc += M_in[worldid, 0, madr] * vec[worldid, col_idof]
+
+  res[worldid, idofid] = acc
+
+
+@wp.kernel
+def _mul_m_island_dense(
+  # Model:
+  nv: int,
+  # Data in:
+  nidof_in: wp.array[int],
+  M_in: wp.array3d[float],
+  map_dof2idof_in: wp.array2d[int],
+  map_idof2dof_in: wp.array2d[int],
+  # In:
+  idof_islandid_in: wp.array2d[int],
+  vec: wp.array2d[float],
+  island_done_in: wp.array2d[bool],
+  check_skip: int,
+  # Out:
+  res: wp.array2d[float],
+):
+  """Dense island mul_m for ALL islands in parallel."""
+  worldid, idofid = wp.tid()
+
+  nidof = nidof_in[worldid]
+  if idofid >= nidof:
+    return
+
+  islandid = idof_islandid_in[worldid, idofid]
+  if islandid < 0:
+    return
+
+  if check_skip:
+    if island_done_in[worldid, islandid]:
+      return
+
+  dof = map_idof2dof_in[worldid, idofid]
+
+  acc = float(0.0)
+  for j in range(nv):
+    col_idof = map_dof2idof_in[worldid, j]
+    # skip unconstrained DOFs
+    if col_idof < nidof:
+      acc += M_in[worldid, dof, j] * vec[worldid, col_idof]
+
+  res[worldid, idofid] = acc
+
+
+@event_scope
+def mul_m_island(
+  m: Model,
+  d: Data,
+  res: wp.array2d[float],
+  vec: wp.array2d[float],
+  nidof: wp.array[int],
+  map_idof2dof: wp.array2d[int],
+  map_dof2idof: wp.array2d[int],
+  idof_islandid: wp.array2d[int],
+  island_done: Optional[wp.array] = None,
+  M: Optional[wp.array] = None,
+):
+  """Multiply island-local vectors by inertia matrix for all islands in parallel.
+
+  Args:
+    m: The model containing kinematic and dynamic information.
+    d: The data object containing the current state and output arrays.
+    res: Result: qM @ vec (island-local DOF order).
+    vec: Input vector (island-local DOF order).
+    nidof: Number of island DOFs per world.
+    map_idof2dof: Island-local DOF → global DOF map.
+    map_dof2idof: Global DOF → island-local DOF map.
+    idof_islandid: Island ID per island-local DOF.
+    island_done: Per-island done flags (nworld, ntree).
+    M: Optional mass matrix override.
+  """
+  check_skip = int(island_done is not None)
+  island_done = island_done or wp.empty((0, 0), dtype=bool)
+
+  if M is None:
+    M = d.M
+
+  if m.is_sparse:
+    wp.launch(
+      _mul_m_island_sparse,
+      dim=(d.nworld, m.nv),
+      inputs=[
+        m.M_mulm_rowadr,
+        m.M_mulm_col,
+        m.M_mulm_madr,
+        nidof,
+        M,
+        map_dof2idof,
+        map_idof2dof,
+        idof_islandid,
+        vec,
+        island_done,
+        check_skip,
+      ],
+      outputs=[res],
+    )
+  else:
+    wp.launch(
+      _mul_m_island_dense,
+      dim=(d.nworld, m.nv),
+      inputs=[
+        m.nv,
+        nidof,
+        M,
+        map_dof2idof,
+        map_idof2dof,
+        idof_islandid,
+        vec,
+        island_done,
+        check_skip,
+      ],
+      outputs=[res],
+    )
+
+
+@cache_kernel
+def _solve_LD_sparse_island(nv: int, nlevels: int):
+  """Sparse backsubstitution with island-local index remapping.
+
+  Same algorithm as _solve_LD_sparse_fused, but reads/writes x/y in
+  island-local DOF order. The L/D factorization stays in global DOF order.
+  """
+
+  @wp.func_native(snippet="WP_TILE_SYNC();")
+  def _syncthreads():
+    pass
+
+  @wp.kernel(module="unique", enable_backward=False)
+  def kernel(
+    # Data in:
+    nidof_in: wp.array[int],
+    map_dof2idof_in: wp.array2d[int],
+    map_idof2dof_in: wp.array2d[int],
+    # In:
+    L: wp.array3d[float],
+    D: wp.array2d[float],
+    all_updates: wp.array[wp.vec3i],
+    level_offsets: wp.array[int],
+    y: wp.array2d[float],
+    # Out:
+    x_out: wp.array2d[float],
+  ):
+    worldid, tid = wp.tid()
+    NLEVELS = wp.static(nlevels)
+    BLOCK_DIM = wp.block_dim()
+    nid = nidof_in[worldid]
+
+    # Copy y to x_out (island-local, only iterate up to nid)
+    for idof in range(tid, nid, BLOCK_DIM):
+      x_out[worldid, idof] = y[worldid, idof]
+    _syncthreads()
+
+    # Forward substitution
+    for level in range(NLEVELS):
+      level_idx = NLEVELS - 1 - level
+      level_offset = level_offsets[level_idx]
+      level_size = level_offsets[level_idx + 1] - level_offset
+
+      for u in range(tid, level_size, BLOCK_DIM):
+        update = all_updates[level_offset + u]
+        i, k, Madr_ki = update[0], update[1], update[2]
+        idof_i = map_dof2idof_in[worldid, i]
+        if idof_i < nid:
+          idof_k = map_dof2idof_in[worldid, k]
+          wp.atomic_sub(x_out[worldid], idof_i, L[worldid, 0, Madr_ki] * x_out[worldid, idof_k])
+      _syncthreads()
+
+    # Diagonal multiply (only iterate up to nid)
+    for idof in range(tid, nid, BLOCK_DIM):
+      dofid = map_idof2dof_in[worldid, idof]
+      x_out[worldid, idof] *= D[worldid, dofid]
+    _syncthreads()
+
+    # Backward substitution
+    for level in range(NLEVELS):
+      level_idx = level
+      level_offset = level_offsets[level_idx]
+      level_size = level_offsets[level_idx + 1] - level_offset
+
+      for u in range(tid, level_size, BLOCK_DIM):
+        update = all_updates[level_offset + u]
+        i, k, Madr_ki = update[0], update[1], update[2]
+        idof_k = map_dof2idof_in[worldid, k]
+        if idof_k < nid:
+          idof_i = map_dof2idof_in[worldid, i]
+          wp.atomic_sub(x_out[worldid], idof_k, L[worldid, 0, Madr_ki] * x_out[worldid, idof_i])
+      _syncthreads()
+
+  return kernel
+
+
+@cache_kernel
+def _tile_cholesky_solve_island(tile):
+  """Dense Cholesky backsubstitution with island-local index remapping.
+
+  L is loaded from global DOF offsets (factorization unchanged).
+  y/x are loaded/stored at island-local DOF offsets via map_dof2idof.
+  """
+
+  @wp.kernel(module="unique", enable_backward=False)
+  def cholesky_solve(
+    # Data in:
+    nidof_in: wp.array[int],
+    map_dof2idof_in: wp.array2d[int],
+    # In:
+    L: wp.array3d[float],
+    y: wp.array2d[float],
+    adr: wp.array[int],
+    # Out:
+    x: wp.array2d[float],
+  ):
+    worldid, nodeid = wp.tid()
+    TILE_SIZE = wp.static(tile.size)
+
+    dofid = adr[nodeid]
+    idofid = map_dof2idof_in[worldid, dofid]
+
+    # Skip unconstrained trees (uniform branch — all threads in block agree)
+    if idofid >= nidof_in[worldid]:
+      return
+
+    # L stays in global order
+    L_tile = wp.tile_load(L[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid))
+    # y and x use island-local offsets
+    y_slice = wp.tile_load(y[worldid], shape=(TILE_SIZE,), offset=(idofid,))
+    x_slice = wp.tile_cholesky_solve(L_tile, y_slice)
+    wp.tile_store(x[worldid], x_slice, offset=(idofid,))
+
+  return cholesky_solve
+
+
+@event_scope
+def solve_m_island(
+  m: Model,
+  d: Data,
+  res: wp.array2d[float],
+  vec: wp.array2d[float],
+  nidof: wp.array[int],
+  map_idof2dof: wp.array2d[int],
+):
+  """Compute res = M^{-1} @ vec for island-local DOFs.
+
+  Args:
+    m: Model.
+    d: Data.
+    res: Output in island-local DOF order.
+    vec: Input in island-local DOF order.
+    nidof: Number of island DOFs per world.
+    map_idof2dof: Island-local DOF -> global DOF map.
+  """
+  if m.is_sparse:
+    nlevels = len(m.qLD_updates)
+    if wp.get_device().is_cuda:
+      dim_block = m.block_dim.solve_LD_sparse_fused
+    else:
+      dim_block = 1
+
+    wp.launch(
+      _solve_LD_sparse_island(m.nv, nlevels),
+      dim=(d.nworld, dim_block),
+      inputs=[
+        d.nidof,
+        d.map_dof2idof,
+        map_idof2dof,
+        d.qLD,
+        d.qLDiagInv,
+        m.qLD_all_updates,
+        m.qLD_level_offsets,
+        vec,
+      ],
+      outputs=[res],
+      block_dim=dim_block,
+    )
+  else:
+    for tile in m.M_tiles:
+      wp.launch_tiled(
+        _tile_cholesky_solve_island(tile),
+        dim=(d.nworld, tile.adr.size),
+        inputs=[d.nidof, d.map_dof2idof, d.qLD, vec, tile.adr],
+        outputs=[res],
+        block_dim=m.block_dim.cholesky_solve,
+      )
+
+
 @wp.kernel
 def _apply_ft(
   # Model:
@@ -604,11 +938,13 @@ def get_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
     nbody: int,
     neq: int,
     nmocap: int,
+    nhistory: int,
     # Data in:
     time_in: wp.array[float],
     qpos_in: wp.array2d[float],
     qvel_in: wp.array2d[float],
     act_in: wp.array2d[float],
+    history_in: wp.array2d[float],
     qacc_warmstart_in: wp.array2d[float],
     ctrl_in: wp.array2d[float],
     qfrc_applied_in: wp.array2d[float],
@@ -647,6 +983,10 @@ def get_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
           for j in range(na):
             state_out[worldid, adr + j] = act_in[worldid, j]
           adr += na
+        elif element == State.HISTORY:
+          for j in range(nhistory):
+            state_out[worldid, adr + j] = history_in[worldid, j]
+          adr += nhistory
         elif element == State.WARMSTART:
           for j in range(nv):
             state_out[worldid, adr + j] = qacc_warmstart_in[worldid, j]
@@ -700,10 +1040,12 @@ def get_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
       m.nbody,
       m.neq,
       m.nmocap,
+      m.nhistory,
       d.time,
       d.qpos,
       d.qvel,
       d.act,
+      d.history,
       d.qacc_warmstart,
       d.ctrl,
       d.qfrc_applied,
@@ -743,6 +1085,7 @@ def set_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
     nbody: int,
     neq: int,
     nmocap: int,
+    nhistory: int,
     # In:
     sig_in: int,
     active_in: wp.array[bool],
@@ -752,6 +1095,7 @@ def set_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
     qpos_out: wp.array2d[float],
     qvel_out: wp.array2d[float],
     act_out: wp.array2d[float],
+    history_out: wp.array2d[float],
     qacc_warmstart_out: wp.array2d[float],
     ctrl_out: wp.array2d[float],
     qfrc_applied_out: wp.array2d[float],
@@ -785,6 +1129,10 @@ def set_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
           for j in range(na):
             act_out[worldid, j] = state_in[worldid, adr + j]
           adr += na
+        elif element == State.HISTORY:
+          for j in range(nhistory):
+            history_out[worldid, j] = state_in[worldid, adr + j]
+          adr += nhistory
         elif element == State.WARMSTART:
           for j in range(nv):
             qacc_warmstart_out[worldid, j] = state_in[worldid, adr + j]
@@ -844,6 +1192,7 @@ def set_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
       m.nbody,
       m.neq,
       m.nmocap,
+      m.nhistory,
       int(sig),
       active or wp.ones(d.nworld, dtype=bool),
       state,
@@ -853,6 +1202,7 @@ def set_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
       d.qpos,
       d.qvel,
       d.act,
+      d.history,
       d.qacc_warmstart,
       d.ctrl,
       d.qfrc_applied,
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py
index 01c974c0..77c6cca2 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py
@@ -17,18 +17,19 @@ import enum
 from typing import Callable
 
 import mujoco
-from mujoco.mjx.third_party.mujoco_warp._src.util_pkg import check_version
 import numpy as np
 import warp as wp
 
+from mujoco.mjx.third_party.mujoco_warp._src.util_pkg import check_version
+
 MJ_MINVAL = mujoco.mjMINVAL
 MJ_MAXVAL = mujoco.mjMAXVAL
 MJ_MINIMP = mujoco.mjMINIMP  # minimum constraint impedance
 MJ_MAXIMP = mujoco.mjMAXIMP  # maximum constraint impedance
 MJ_MAXCONPAIR = mujoco.mjMAXCONPAIR
 MJ_MINMU = mujoco.mjMINMU  # minimum friction
-# True if MuJoCo >= 3.9.0 (new margin/gap semantics: includemargin = margin)
-_NEW_GAP_SEMANTICS = check_version("mujoco>=3.9.0")
+NEW_GAP_SEMANTICS = check_version("mujoco>=3.9.0.dev914519929")
+TACTILE_DEPTH_SEMANTICS = check_version("mujoco>=3.9.0.dev921980899")
 # maximum size (by number of edges) of an horizon in EPA algorithm
 MJ_MAX_EPAHORIZON = 24
 # maximum average number of trianglarfaces EPA can insert at each iteration
@@ -47,6 +48,25 @@ class BlockDim:
   """Block dimension 'block_dim' settings for wp.launch_tiled.
 
   TODO(team): experimental and may be removed
+
+  Attributes:
+    segmented_sort: segmented sort block dimension (collision_driver)
+    euler_dense: Euler dense block dimension (forward)
+    actuator_velocity: actuator velocity block dimension (forward)
+    ray: ray block dimension (ray)
+    contact_sort: contact sort block dimension (sensor)
+    energy_vel_kinetic: energy velocity kinetic block dimension (sensor)
+    cholesky_factorize: Cholesky factorize block dimension (smooth)
+    cholesky_solve: Cholesky solve block dimension (smooth)
+    cholesky_factorize_solve: Cholesky factorize and solve block dimension (smooth)
+    solve_LD_sparse_fused: solve LD sparse fused block dimension (smooth)
+    update_gradient_cholesky: update gradient Cholesky block dimension (solver)
+    update_gradient_cholesky_blocked: update gradient Cholesky blocked block dimension (solver)
+    update_gradient_JTDAJ_sparse: update gradient JTDAJ sparse block dimension (solver)
+    update_gradient_JTDAJ_dense: update gradient JTDAJ dense block dimension (solver)
+    linesearch_iterative: linesearch iterative block dimension (solver)
+    contact_jac_tiled: contact Jacobian tiled block dimension (solver)
+    qderiv_actuator_dense: qderiv actuator dense block dimension (derivative)
   """
 
   # collision_driver
@@ -131,13 +151,8 @@ class ProjectionType(enum.IntEnum):
     ORTHOGRAPHIC: orthographic projection
   """
 
-  # TODO(team): remove after mjwarp depends on mujoco > 3.4.0 in pyproject.toml
-  if hasattr(mujoco, "mjtProjection"):
-    PERSPECTIVE = mujoco.mjtProjection.mjPROJ_PERSPECTIVE
-    ORTHOGRAPHIC = mujoco.mjtProjection.mjPROJ_ORTHOGRAPHIC
-  else:
-    PERSPECTIVE = 0
-    ORTHOGRAPHIC = 1
+  PERSPECTIVE = mujoco.mjtProjection.mjPROJ_PERSPECTIVE
+  ORTHOGRAPHIC = mujoco.mjtProjection.mjPROJ_ORTHOGRAPHIC
 
 
 class Stage(enum.IntEnum):
@@ -188,6 +203,7 @@ class DisableBit(enum.IntFlag):
     EULERDAMP:    implicit damping for Euler integration
     NATIVECCD:    native convex collision detection (ignored in MJWarp)
     ISLAND:       constraint islands
+    MULTICCD:     disable multiple contacts with CCD
   """
 
   CONSTRAINT = mujoco.mjtDisableBit.mjDSBL_CONSTRAINT
@@ -217,7 +233,6 @@ class EnableBit(enum.IntFlag):
   Attributes:
     ENERGY: energy computation
     INVDISCRETE: discrete-time inverse dynamics
-    MULTICCD: multiple contacts with CCD
   """
 
   ENERGY = mujoco.mjtEnableBit.mjENBL_ENERGY
@@ -338,12 +353,13 @@ class IntegratorType(enum.IntEnum):
     EULER: semi-implicit Euler
     RK4: 4th-order Runge Kutta
     IMPLICITFAST: implicit in velocity, no rne derivative
+    IMPLICIT: implicit in velocity, with rne derivative
   """
 
   EULER = mujoco.mjtIntegrator.mjINT_EULER
   RK4 = mujoco.mjtIntegrator.mjINT_RK4
   IMPLICITFAST = mujoco.mjtIntegrator.mjINT_IMPLICITFAST
-  # unsupported: IMPLICIT
+  IMPLICIT = mujoco.mjtIntegrator.mjINT_IMPLICIT
 
 
 class GeomType(enum.IntEnum):
@@ -617,6 +633,7 @@ class State(enum.IntEnum):
     QPOS: position
     QVEL: velocity
     ACT: actuator activation
+    HISTORY: delay/interval history buffers
     WARMSTART: acceleration used for warmstart
     CTRL: control
     QFRC_APPLIED: applied generalized force
@@ -625,7 +642,7 @@ class State(enum.IntEnum):
     MOCAP_POS: positions of mocap bodies
     MOCAP_QUAT: orientations of mocap bodies
     NSTATE: number of state elements
-    PHYSICS: QPOS | QVEL | ACT
+    PHYSICS: TIME | QPOS | QVEL | ACT | HISTORY
     FULLPHYSICS: TIME | PHYSICS | PLUGIN
     USER: CTRL | QFRC_APPLIED | XFRC_APPLIED | EQ_ACTIVE | MOCAP_POS | MOCAP_QUAT | USERDATA
     INTEGRATION: FULLPHYSICS | USER | WARMSTART
@@ -635,6 +652,7 @@ class State(enum.IntEnum):
   QPOS = mujoco.mjtState.mjSTATE_QPOS
   QVEL = mujoco.mjtState.mjSTATE_QVEL
   ACT = mujoco.mjtState.mjSTATE_ACT
+  HISTORY = mujoco.mjtState.mjSTATE_HISTORY
   WARMSTART = mujoco.mjtState.mjSTATE_WARMSTART
   CTRL = mujoco.mjtState.mjSTATE_CTRL
   QFRC_APPLIED = mujoco.mjtState.mjSTATE_QFRC_APPLIED
@@ -643,7 +661,7 @@ class State(enum.IntEnum):
   MOCAP_POS = mujoco.mjtState.mjSTATE_MOCAP_POS
   MOCAP_QUAT = mujoco.mjtState.mjSTATE_MOCAP_QUAT
   NSTATE = mujoco.mjtState.mjNSTATE
-  PHYSICS = mujoco.mjtState.mjSTATE_PHYSICS
+  PHYSICS = mujoco.mjtState.mjSTATE_PHYSICS  # includes HISTORY
   FULLPHYSICS = mujoco.mjtState.mjSTATE_FULLPHYSICS
   USER = mujoco.mjtState.mjSTATE_USER
   INTEGRATION = mujoco.mjtState.mjSTATE_INTEGRATION
@@ -861,6 +879,7 @@ class Model:
     ntree: number of kinematic trees
     nM: number of non-zeros in sparse inertia matrix
     nC: number of non-zeros in sparse body-dof matrix
+    nD: number of non-zeros in sparse derivative matrix
     ngeom: number of geoms
     nsite: number of sites
     ncam: number of cameras
@@ -898,6 +917,7 @@ class Model:
     nJmom: number of non-zeros in actuator_moment
     ngravcomp: number of bodies with nonzero gravcomp
     nsensordata: number of elements in sensor data vector
+    nhistory: number of history buffer entries
     opt: physics options
     stat: model statistics
     qpos0: qpos values at default pose                       (*, nq)
@@ -1129,6 +1149,9 @@ class Model:
     actuator_trnid: transmission id: joint, tendon, site     (nu, 2)
     actuator_actadr: first activation address; -1: stateless (nu,)
     actuator_actnum: number of activation variables          (nu,)
+    actuator_history: history buffer sizes                   (nu, 2)
+    actuator_historyadr: history buffer address              (nu,)
+    actuator_delay: delay in seconds                         (nu,)
     actuator_ctrllimited: is control limited                 (nu,)
     actuator_forcelimited: is force limited                  (nu,)
     actuator_actlimited: is activation limited               (nu,)
@@ -1153,12 +1176,22 @@ class Model:
     sensor_dim: number of scalar outputs                     (nsensor,)
     sensor_adr: address in sensor array                      (nsensor,)
     sensor_cutoff: cutoff for real and positive; 0: ignore   (nsensor,)
+    sensor_history: history buffer sizes                     (nsensor, 2)
+    sensor_historyadr: history buffer address                (nsensor,)
+    sensor_delay: delay in seconds                           (nsensor,)
+    sensor_interval: sensor interval and phase               (nsensor, 2)
     plugin: globally registered plugin slot number           (nplugin,)
     plugin_attr: config attributes of geom plugin            (nplugin, _NPLUGINATTR)
-    M_rownnz: number of non-zeros in each row of qM          (nv,)
-    M_rowadr: index of each row in qM                        (nv,)
-    M_colind: column indices of non-zeros in qM              (nM,)
+    M_rownnz: number of non-zeros in each row of M           (nv,)
+    M_rowadr: index of each row in M                         (nv,)
+    M_colind: column indices of non-zeros in M               (nC,)
     mapM2M: index mapping from M (legacy) to M (CSR)         (nC)
+    D_rownnz: non-zeros per row in D-structure               (nv,)
+    D_rowadr: row start addresses in D-structure             (nv,)
+    D_diag: diagonal element index within each row           (nv,)
+    D_colind: column indices in D-structure                  (nD,)
+    mapM2D: index mapping from M to D                        (nD,)
+    mapD2M: index mapping from D to M                        (nC,)
     flex_vertflexid: flex id for each flex vertex            (nflexvert,)
 
   warp only fields:
@@ -1187,8 +1220,8 @@ class Model:
     jnt_limited_slide_hinge_adr: limited/slide/hinge jntadr
     jnt_limited_ball_adr: limited/ball jntadr
     body_isdofancestor: precomputed mask of which DOFs affect each body
-    dof_tri_row: dof lower triangle row (used in solver)
-    dof_tri_col: dof lower triangle col (used in solver)
+    dof_tri_row: dof upper triangle row (used in solver)
+    dof_tri_col: dof upper triangle col (used in solver)
     nxn_geom_pair: collision pair geom ids [-2, ngeom-1]
     nxn_geom_pair_filtered: valid collision pair geom ids
                             [-1, ngeom - 1]
@@ -1243,15 +1276,21 @@ class Model:
     sensor_rangefinder_bodyid: bodyid for rangefinder        (nrangefinder,)
     taxel_vertadr: tactile sensor vertex address             (nsensortaxel,)
     taxel_sensorid: address for tactile sensors
-    qM_tiles: tiling configuration
+    M_tiles: tiling configuration
     qLD_updates: tuple of index triples for sparse factorization
     qLD_all_updates: tuple of all levels concatenated
     qLD_level_offsets: tuple of start offsets for each level
-    qM_fullm_i: sparse mass matrix addressing
-    qM_fullm_j: sparse mass matrix addressing
-    qM_mulm_rowadr: sparse matmul row pointers
-    qM_mulm_col: sparse matmul column indices
-    qM_mulm_madr: sparse matmul matrix addresses
+    M_fullm_i: sparse mass matrix addressing
+    M_fullm_j: sparse mass matrix addressing
+    M_elemid: (row, col) -> CSR madr addresses; -1 if not a chain ancestor
+    M_fullm_upper_i: upper-triangle row indices for solver h seeding
+    M_fullm_upper_j: upper-triangle column indices for solver h seeding
+    M_fullm_upper_elemid: source elemid into M_fullm_i/M_fullm_j
+    qD_fullm_i: D-structure row indices for RNE derivatives
+    qD_fullm_j: D-structure column indices for RNE derivatives
+    M_mulm_rowadr: sparse matmul row pointers
+    M_mulm_col: sparse matmul column indices
+    M_mulm_madr: sparse matmul matrix addresses
   """
 
   nq: int
@@ -1264,6 +1303,7 @@ class Model:
   ntree: int
   nM: int
   nC: int
+  nD: int
   ngeom: int
   nsite: int
   ncam: int
@@ -1301,6 +1341,7 @@ class Model:
   nJmom: int
   ngravcomp: int
   nsensordata: int
+  nhistory: int
   opt: Option
   stat: Statistic
   qpos0: array("*", "nq", float)
@@ -1532,6 +1573,9 @@ class Model:
   actuator_trnid: array("nu", wp.vec2i)
   actuator_actadr: array("nu", int)
   actuator_actnum: array("nu", int)
+  actuator_history: array("nu", wp.vec2i)
+  actuator_historyadr: array("nu", int)
+  actuator_delay: array("nu", float)
   actuator_ctrllimited: array("nu", bool)
   actuator_forcelimited: array("nu", bool)
   actuator_actlimited: array("nu", bool)
@@ -1556,12 +1600,22 @@ class Model:
   sensor_dim: array("nsensor", int)
   sensor_adr: array("nsensor", int)
   sensor_cutoff: array("nsensor", float)
+  sensor_history: array("nsensor", wp.vec2i)
+  sensor_historyadr: array("nsensor", int)
+  sensor_delay: array("nsensor", float)
+  sensor_interval: array("nsensor", wp.vec2)
   plugin: array("nplugin", int)
   plugin_attr: array("nplugin", vec_pluginattr)
   M_rownnz: array("nv", int)
   M_rowadr: array("nv", int)
   M_colind: array("nC", int)
   mapM2M: array("nC", int)
+  D_rownnz: array("nv", int)
+  D_rowadr: array("nv", int)
+  D_diag: array("nv", int)
+  D_colind: array("nD", int)
+  mapM2D: array("nD", int)
+  mapD2M: array("nC", int)
   flex_vertflexid: array("nflexvert", int)
   # warp only fields:
   callback: Callback
@@ -1635,23 +1689,32 @@ class Model:
   sensor_rangefinder_bodyid: array("nrangefinder", int)
   taxel_vertadr: array("nsensortaxel", int)
   taxel_sensorid: wp.array[int]
-  qM_tiles: tuple[TileSet, ...]
+  M_tiles: tuple[TileSet, ...]
   qLD_updates: tuple[wp.array[wp.vec3i], ...]
   qLD_all_updates: wp.array[wp.vec3i]
   qLD_level_offsets: wp.array[int]
-  qM_fullm_i: wp.array[int]
-  qM_fullm_j: wp.array[int]
+  # TODO(team): Remove M_fullm_i/j and M_elemid by iterating the M CSR layout
+  # directly in the solver/derivative kernels
+  M_fullm_i: wp.array[int]
+  M_fullm_j: wp.array[int]
+  M_elemid: wp.array2d[int]  # (row, col) -> CSR madr address; -1 if col is not a chain ancestor of row
+  M_fullm_upper_i: wp.array[int]
+  M_fullm_upper_j: wp.array[int]
+  M_fullm_upper_elemid: wp.array[int]
+  qD_fullm_i: wp.array[int]  # D-structure (full square) row indices for RNE derivatives
+  qD_fullm_j: wp.array[int]  # D-structure (full square) column indices for RNE derivatives
   # Gather-based sparse mul_m indices (thread per DOF, no atomics)
-  qM_mulm_rowadr: wp.array[int]  # start address for each row [nv+1]
-  qM_mulm_col: wp.array[int]  # column index to gather from
-  qM_mulm_madr: wp.array[int]  # matrix address to read
+  M_mulm_rowadr: wp.array[int]  # start address for each row [nv+1]
+  M_mulm_col: wp.array[int]  # column index to gather from
+  M_mulm_madr: wp.array[int]  # matrix address to read
 
 
 class ContactType(enum.IntFlag):
   """Type of contact.
 
-  CONSTRAINT: contact for constraint solver.
-  SENSOR: contact for collision sensor (GEOMDIST, GEOMNORMAL, GEOMFROMTO).
+  Attributes:
+    CONSTRAINT: contact for constraint solver
+    SENSOR: contact for collision sensor (GEOMDIST, GEOMNORMAL, GEOMFROMTO)
   """
 
   CONSTRAINT = 1
@@ -1673,6 +1736,8 @@ class Contact:
     solimp: constraint solver impedance                              (naconmax, 5)
     dim: contact space dimensionality: 1, 3, 4 or 6                  (naconmax,)
     geom: geom ids; -1 for flex                                      (naconmax, 2)
+    flex: flex ids; -1 for geom                                      (naconmax, 2)
+    vert: vertex ids for flex/mesh contact                           (naconmax, 2)
     efc_address: address in efc; -1: not included                    (naconmax, nmaxpyramid)
     worldid: world id                                                (naconmax,)
     type: ContactType                                                (naconmax,)
@@ -1711,9 +1776,9 @@ class Constraint:
     J_rowadr: row start address in colind array       (nworld, 0) dense
                                                       (nworld, njmax) sparse
     J_colind: column indices in J                     (nworld, 0, 0) dense
-                                                      (nworld, 1, njmax * nv) sparse
+                                                      (nworld, 1, njmax_nnz) sparse
     J: constraint Jacobian                            (nworld, njmax_pad, nv_pad) dense
-                                                      (nworld, 1, njmax * nv) sparse
+                                                      (nworld, 1, njmax_nnz) sparse
     pos: constraint position (equality, contact)      (nworld, njmax)
     margin: inclusion margin (contact)                (nworld, njmax)
     D: constraint mass                                (nworld, njmax_pad)
@@ -1722,6 +1787,20 @@ class Constraint:
     frictionloss: frictionloss (friction)             (nworld, njmax)
     force: constraint force in constraint space       (nworld, njmax)
     state: constraint state                           (nworld, njmax_pad)
+    island: island ID per constraint                  (nworld, njmax)
+    itype: island constraint type                     (nworld, njmax)
+    iid: island constraint id                         (nworld, njmax)
+    iJ_rownnz: island J_rownnz                        (nworld, njmax)
+    iJ_rowadr: island J_rowadr                        (nworld, njmax)
+    iJ_colind: island J_colind                        (nworld, 0, 0) dense
+                                                      (nworld, 1, njmax_nnz) sparse
+    iJ: island J                                      (nworld, njmax, nv) dense
+                                                      (nworld, 1, njmax_nnz) sparse
+    iD: island constraint mass                        (nworld, njmax_pad)
+    iaref: island aref                                (nworld, njmax)
+    ifrictionloss: island frictionloss                (nworld, njmax)
+    iforce: island force                              (nworld, njmax)
+    istate: island state                              (nworld, njmax_pad)
   warp only fields:
     Ma: M*qacc                                        (nworld, nv)
     Jqvel: J*qvel                                     (nworld, njmax)
@@ -1729,8 +1808,8 @@ class Constraint:
 
   type: array("nworld", "njmax", int)
   id: array("nworld", "njmax", int)
-  J_rownnz: wp.array2d[int]
-  J_rowadr: wp.array2d[int]
+  J_rownnz: array("nworld", "njmax", int)
+  J_rowadr: array("nworld", "njmax", int)
   J_colind: wp.array3d[int]
   J: wp.array3d[float]
   pos: array("nworld", "njmax", float)
@@ -1741,9 +1820,22 @@ class Constraint:
   frictionloss: array("nworld", "njmax", float)
   force: array("nworld", "njmax", float)
   state: array("nworld", "njmax_pad", int)
+  island: array("nworld", "njmax", int)
   Ma: array("nworld", "nv", float)
   Jqvel: array("nworld", "njmax", float)
 
+  itype: array("nworld", "njmax", int)
+  iid: array("nworld", "njmax", int)
+  iJ_rownnz: array("nworld", "njmax", int)
+  iJ_rowadr: array("nworld", "njmax", int)
+  iJ_colind: wp.array3d[int]
+  iJ: wp.array3d[float]
+  iD: array("nworld", "njmax_pad", float)
+  iaref: array("nworld", "njmax", float)
+  ifrictionloss: array("nworld", "njmax", float)
+  iforce: array("nworld", "njmax", float)
+  istate: array("nworld", "njmax_pad", int)
+
 
 @dataclasses.dataclass
 class Data:
@@ -1756,11 +1848,13 @@ class Data:
     nl: number of limit constraints                             (nworld,)
     nefc: number of constraints                                 (nworld,)
     nisland: number of constraint islands                       (nworld,)
+    nidof: total DOFs in islands                                (nworld,)
     time: simulation time                                       (nworld,)
     energy: potential, kinetic energy                           (nworld, 2)
     qpos: position                                              (nworld, nq)
     qvel: velocity                                              (nworld, nv)
     act: actuator activation                                    (nworld, na)
+    history: history buffer for delays                          (nworld, nhistory)
     qacc_warmstart: acceleration used for warmstart             (nworld, nv)
     ctrl: control                                               (nworld, nu)
     qfrc_applied: applied generalized force                     (nworld, nv)
@@ -1791,7 +1885,7 @@ class Data:
     cinert: com-based body inertia and mass                     (nworld, nbody, 10)
     flexvert_xpos: cartesian flex vertex positions              (nworld, nflexvert, 3)
     flexedge_J: edge length Jacobian                            (nworld, nJfe)
-    flexedge_length: flex edge lengths                          (nworld, nflexedge, 1)
+    flexedge_length: flex edge lengths                          (nworld, nflexedge)
     ten_wrapadr: start address of tendon's path                 (nworld, ntendon)
     ten_wrapnum: number of wrap points in path                  (nworld, ntendon)
     ten_J: tendon Jacobian                                      (nworld, nJten)
@@ -1804,10 +1898,10 @@ class Data:
     moment_colind: column indices in sparse actuator_moment     (nworld, nJmom)
     actuator_moment: actuator moments                           (nworld, nJmom)
     crb: com-based composite inertia and mass                   (nworld, nbody, 10)
-    qM: total inertia                                           (nworld, nv, nv) if dense
-                                                                (nworld, 1, nM) if sparse
-    qLD: L'*D*L factorization of M                              (nworld, nv, nv) if dense
+    M: total inertia                                            (nworld, nv_pad, nv_pad) if dense
                                                                 (nworld, 1, nC) if sparse
+    qLD: upper Cholesky factorization                           (nworld, nv, nv) if dense
+         L'*D*L factorization of M                              (nworld, 1, nC) if sparse
     qLDiagInv: 1/diag(D)                                        (nworld, nv)
     flexedge_velocity: flex edge velocities                     (nworld, nflexedge)
     ten_velocity: tendon velocities                             (nworld, ntendon)
@@ -1822,6 +1916,7 @@ class Data:
     qfrc_passive: total passive force                           (nworld, nv)
     subtree_linvel: linear velocity of subtree com              (nworld, nbody, 3)
     subtree_angmom: angular momentum about subtree com          (nworld, nbody, 3)
+    qLU: sparse LU factorization of (M - dt*qDeriv)             (nworld, 1, nD)
     actuator_force: actuator force in actuation space           (nworld, nu)
     qfrc_actuator: actuator force                               (nworld, nv)
     qfrc_smooth: net unconstrained force                        (nworld, nv)
@@ -1836,6 +1931,23 @@ class Data:
     contact: contact data
     efc: constraint data
     tree_island: island ID per tree (-1 if unconstrained)       (nworld, ntree)
+    dof_island: island ID per DOF (-1 if unconstrained)         (nworld, nv)
+    island_dofadr: island start address in dof vector           (nworld, ntree)
+    island_nv: DOFs per island                                  (nworld, ntree)
+    island_nefc: constraints per island                         (nworld, ntree)
+    island_ne: equality constraints per island                  (nworld, ntree)
+    island_nf: friction constraints per island                  (nworld, ntree)
+    island_efcadr: island start address in efc vector           (nworld, ntree)
+    map_dof2idof: global DOF -> island-local DOF                (nworld, nv)
+    map_idof2dof: island-local DOF -> global DOF                (nworld, nv)
+    map_efc2iefc: global EFC -> island-local EFC                (nworld, njmax)
+    map_iefc2efc: island-local EFC -> global EFC                (nworld, njmax)
+    dof_islandid: island ID per island-DOF                      (nworld, nv)
+    efc_islandid: island ID per island-EFC                      (nworld, njmax)
+    iqacc: island-local qacc                                    (nworld, nv)
+    iqacc_smooth: island-local qacc_smooth                      (nworld, nv)
+    iqfrc_smooth: island-local qfrc_smooth                      (nworld, nv)
+    iqfrc_constraint: island-local qfrc_constraint              (nworld, nv)
 
   warp only fields:
     nworld: number of worlds
@@ -1854,11 +1966,13 @@ class Data:
   nl: array("nworld", int)
   nefc: array("nworld", int)
   nisland: array("nworld", int)
+  nidof: array("nworld", int)
   time: array("nworld", float)
   energy: array("nworld", wp.vec2)
   qpos: array("nworld", "nq", float)
   qvel: array("nworld", "nv", float)
   act: array("nworld", "na", float)
+  history: array("nworld", "nhistory", float)
   qacc_warmstart: array("nworld", "nv", float)
   ctrl: array("nworld", "nu", float)
   qfrc_applied: array("nworld", "nv", float)
@@ -1902,7 +2016,7 @@ class Data:
   moment_colind: array("nworld", "nJmom", int)
   actuator_moment: array("nworld", "nJmom", float)
   crb: array("nworld", "nbody", vec10)
-  qM: wp.array3d[float]
+  M: wp.array3d[float]
   qLD: wp.array3d[float]
   qLDiagInv: array("nworld", "nv", float)
   flexedge_velocity: array("nworld", "nflexedge", float)
@@ -1918,6 +2032,7 @@ class Data:
   qfrc_passive: array("nworld", "nv", float)
   subtree_linvel: array("nworld", "nbody", wp.vec3)
   subtree_angmom: array("nworld", "nbody", wp.vec3)
+  qLU: array("nworld", 1, "nD", float)
   actuator_force: array("nworld", "nu", float)
   qfrc_actuator: array("nworld", "nv", float)
   qfrc_smooth: array("nworld", "nv", float)
@@ -1930,6 +2045,23 @@ class Data:
   contact: Contact
   efc: Constraint
   tree_island: array("nworld", "ntree", int)
+  dof_island: array("nworld", "nv", int)
+  island_dofadr: array("nworld", "ntree", int)
+  island_nv: array("nworld", "ntree", int)
+  island_nefc: array("nworld", "ntree", int)
+  island_ne: array("nworld", "ntree", int)
+  island_nf: array("nworld", "ntree", int)
+  island_efcadr: array("nworld", "ntree", int)
+  map_dof2idof: array("nworld", "nv", int)
+  map_idof2dof: array("nworld", "nv", int)
+  map_efc2iefc: array("nworld", "njmax", int)
+  map_iefc2efc: array("nworld", "njmax", int)
+  dof_islandid: array("nworld", "nv", int)
+  efc_islandid: array("nworld", "njmax", int)
+  iqacc: wp.array2d[float]
+  iqacc_smooth: wp.array2d[float]
+  iqfrc_smooth: wp.array2d[float]
+  iqfrc_constraint: wp.array2d[float]
 
   # warp only fields:
   nworld: int
@@ -1942,6 +2074,77 @@ class Data:
   ncollision: array(1, int)
 
 
+@dataclasses.dataclass
+class InverseContext:
+  """Workspace arrays for inverse dynamics."""
+
+  Jaref: wp.array2d[float]
+  search_dot: wp.array[float]
+  gauss: wp.array[float]
+  cost: wp.array[float]
+  prev_cost: wp.array[float]
+  done: wp.array[bool]
+  changed_efc_ids: wp.array2d[int]
+  changed_efc_count: wp.array[int]
+
+
+@dataclasses.dataclass
+class IslandSolverContext:
+  """Workspace arrays for island constraint solver."""
+
+  # Re-ordered workspace arrays (sized per-dof / per-constraint)
+  Jaref: wp.array2d[float]
+  jv: wp.array2d[float]
+  search: wp.array2d[float]
+  mv: wp.array2d[float]
+  grad: wp.array2d[float]
+  Mgrad: wp.array2d[float]
+  prev_grad: wp.array2d[float]
+  prev_Mgrad: wp.array2d[float]
+  h: wp.array3d[float]
+
+  # Per-island solver scalars (nworld, ntree)
+  cost: wp.array2d[float]
+  prev_cost: wp.array2d[float]
+  gauss: wp.array2d[float]
+  search_dot: wp.array2d[float]
+  grad_dot: wp.array2d[float]
+  done: wp.array2d[bool]  # per-island convergence
+  solver_niter: wp.array2d[int]  # iterations per island
+  beta: wp.array2d[float]
+  alpha: wp.array2d[float]
+  Ma: wp.array2d[float]  # island-local Ma (nworld, nv)
+
+
+@dataclasses.dataclass
+class SolverContext:
+  """Workspace arrays for constraint solver."""
+
+  Jaref: wp.array2d[float]
+  search_dot: wp.array[float]
+  gauss: wp.array[float]
+  cost: wp.array[float]
+  prev_cost: wp.array[float]
+  done: wp.array[bool]
+  grad: wp.array2d[float]
+  grad_dot: wp.array[float]
+  Mgrad: wp.array2d[float]
+  search: wp.array2d[float]
+  mv: wp.array2d[float]
+  jv: wp.array2d[float]
+  quad: wp.array2d[wp.vec3]
+  quad_gauss: wp.array[wp.vec3]
+  alpha: wp.array[float]
+  prev_grad: wp.array2d[float]
+  prev_Mgrad: wp.array2d[float]
+  beta: wp.array[float]
+  h: wp.array3d[float]
+  hfactor: wp.array3d[float]
+  # Incremental Hessian update (Newton only)
+  changed_efc_ids: wp.array2d[int]
+  changed_efc_count: wp.array[int]
+
+
 @dataclasses.dataclass
 class RenderContext:
   """Context for rendering.
@@ -1952,6 +2155,8 @@ class RenderContext:
     cam_id_map: camera id map
     use_textures: whether to use textures
     use_shadows: whether to use shadows
+    use_ambient_lighting: whether to use ambient lighting
+    background_color: background color
     use_precomputed_rays: whether to use precomputed rays
     bvh_ngeom: number of geometries in the BVH
     enabled_geom_ids: enabled geometry ids
@@ -1971,7 +2176,10 @@ class RenderContext:
     flex_bvh_id: per-flex BVH ids
     flex_group_root: per-flex group roots (nworld x n_flex_bvh)
     flex_render_smooth: whether to render flex meshes smoothly
-    flex_dim: flex dimension per flex (1D/2D/3D)
+    bvh_nflexgeom: number of flex geometries in the BVH
+    flex_dim_np: flex dimension per flex (1D/2D/3D)
+    flex_geom_flexid: map from flex geom ID to flex ID
+    flex_geom_edgeid: map from flex geom ID to flex edge ID
     bvh: scene BVH
     bvh_id: scene BVH id
     lower: lower bounds
@@ -1993,6 +2201,11 @@ class RenderContext:
     render_skybox: whether to shade missed rays with the MuJoCo skybox texture
     skybox_tex_id: index into textures of the skybox (MuJoCo tex_type == SKYBOX), -1 if none
     skybox_face_width: pixel width of one skybox cube face (0 if no skybox)
+    enable_backface_culling: drop primitive ray hits whose normal faces away
+      from the ray (i.e. the ray origin is inside the geom). Matches MuJoCo's
+      mesh-ray rule. When False, the renderer reports inner-surface hits, which
+      is faster but causes a camera placed inside a geom to render that geom's
+      back wall.
   """
 
   nrender: int
@@ -2000,6 +2213,7 @@ class RenderContext:
   cam_id_map: array("ncam", int)
   use_textures: bool
   use_shadows: bool
+  use_ambient_lighting: bool
   background_color: wp.uint32
   use_precomputed_rays: bool
   render_skybox: bool
@@ -2045,3 +2259,4 @@ class RenderContext:
   render_seg: array("ncam", bool)
   znear: float
   total_rays: int
+  enable_backface_culling: bool
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/util_pkg.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/util_pkg.py
index 9510b8fd..d5860af2 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/util_pkg.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/util_pkg.py
@@ -23,7 +23,7 @@ def _parse_version(version_str: str) -> tuple[tuple[int, int | str], ...]:
   """Parse a version string into comparable components.
 
   Dot-separated components form the version. Hyphen-separated suffixes (e.g.,
-  "-foo3") are treated as local build identifiers and stripped before
+  "-newton") are treated as local build identifiers and stripped before
   parsing. Each component is wrapped in a tuple: (0, int) for numeric parts,
   (-1, str) for non-numeric. A (0, 0) sentinel is appended so that stable
   releases sort above pre-release suffixes during Python tuple comparison
@@ -32,13 +32,13 @@ def _parse_version(version_str: str) -> tuple[tuple[int, int | str], ...]:
 
   Args:
     version_str: Version string like "3.5.0", "3.5.0.dev869102767", or
-      "3.9.0-foo3".
+      "3.9.0-newton".
 
   Returns:
     Tuple of (type_order, value) pairs for comparison, where type_order is 0
     for integers and -1 for strings, followed by a (0, 0) sentinel.
   """
-  # Strip local build identifier (e.g., "3.9.0-foo3" -> "3.9.0")
+  # Strip local build identifier (e.g., "3.9.0-newton" -> "3.9.0")
   version_str = version_str.split("-", 1)[0]
   parts = version_str.split(".")
   return tuple([(0, int(p)) if p.isdigit() else (-1, p) for p in parts] + [(0, 0)])
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/pyproject.toml b/mjx/mujoco/mjx/third_party/mujoco_warp/pyproject.toml
index 92750398..eac2e265 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/pyproject.toml
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name="mujoco-warp"
-version = "3.8.0.1"
+version = "3.9.0.1"
 # TODO(team): create a distribution list
 authors = [
     {name = "Newton Developers", email = "mujoco@deepmind.com"},
@@ -31,7 +31,7 @@ dependencies = [
     "etils[epath]",
     "mujoco>=3.8.0",
     "numpy",
-    "warp-lang>=1.12",
+    "warp-lang>=1.13",
 ]
 
 [[tool.uv.index]]
diff --git a/mjx/mujoco/mjx/warp/__init__.py b/mjx/mujoco/mjx/warp/__init__.py
index cc582e4a..28ebd35f 100644
--- a/mjx/mujoco/mjx/warp/__init__.py
+++ b/mjx/mujoco/mjx/warp/__init__.py
@@ -23,6 +23,7 @@ if not typing.TYPE_CHECKING:
   warp: Any = None
   mujoco_warp: Any = None
   mjwp_types: Any = None
+  mjwp_io: Any = None
   WARP_INSTALLED: bool = False
 
   # pylint: disable=g-import-not-at-top
@@ -35,6 +36,7 @@ if not typing.TYPE_CHECKING:
   try:
     from mujoco.mjx.third_party import mujoco_warp
     from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
+    from mujoco.mjx.third_party.mujoco_warp._src import io as mjwp_io
   except (ImportError, RuntimeError) as e:
     print('Failed to import mujoco_warp:', e)
     pass
@@ -67,7 +69,11 @@ else:
     def BlockDim(self, *args, **kwargs):  # pylint: disable=invalid-name
       pass
 
+  class _MjwpIoStub:
+    ENABLE_ISLANDS = True
+
   WARP_INSTALLED: bool = True
   warp: Any = _WpStub()
   mujoco_warp: Any = _MjwpStub()
   mjwp_types: Any = _MjwpTypesStub()
+  mjwp_io: Any = _MjwpIoStub()
diff --git a/mjx/mujoco/mjx/warp/forward.py b/mjx/mujoco/mjx/warp/forward.py
index 52a84999..8bae7621 100644
--- a/mjx/mujoco/mjx/warp/forward.py
+++ b/mjx/mujoco/mjx/warp/forward.py
@@ -53,8 +53,18 @@ _cb = mjwp_types.Callback(
 def _forward_shim(
     # Model
     nworld: int,
+    M_elemid: wp.array2d[int],
+    M_fullm_i: wp.array[int],
+    M_fullm_j: wp.array[int],
+    M_fullm_upper_elemid: wp.array[int],
+    M_fullm_upper_i: wp.array[int],
+    M_fullm_upper_j: wp.array[int],
+    M_mulm_col: wp.array[int],
+    M_mulm_madr: wp.array[int],
+    M_mulm_rowadr: wp.array[int],
     M_rowadr: wp.array[int],
     M_rownnz: wp.array[int],
+    M_tiles: tuple[mjwp_types.TileSet, ...],
     actuator_acc0: wp.array2d[float],
     actuator_actadr: wp.array[int],
     actuator_actearly: wp.array[bool],
@@ -66,6 +76,7 @@ def _forward_shim(
     actuator_cranklength: wp.array2d[float],
     actuator_ctrllimited: wp.array[bool],
     actuator_ctrlrange: wp.array2d[wp.vec2],
+    actuator_delay: wp.array[float],
     actuator_dynprm: wp.array2d[mjwp_types.vec10f],
     actuator_dyntype: wp.array[int],
     actuator_forcelimited: wp.array[bool],
@@ -73,6 +84,8 @@ def _forward_shim(
     actuator_gainprm: wp.array2d[mjwp_types.vec10f],
     actuator_gaintype: wp.array[int],
     actuator_gear: wp.array2d[wp.spatial_vector],
+    actuator_history: wp.array[wp.vec2i],
+    actuator_historyadr: wp.array[int],
     actuator_lengthrange: wp.array2d[wp.vec2],
     actuator_trnid: wp.array[wp.vec2i],
     actuator_trntype: wp.array[int],
@@ -115,7 +128,6 @@ def _forward_shim(
     cam_resolution: wp.array[wp.vec2i],
     cam_sensorsize: wp.array[wp.vec2],
     cam_targetbodyid: wp.array[int],
-    dof_Madr: wp.array[int],
     dof_armature: wp.array2d[float],
     dof_bodyid: wp.array[int],
     dof_damping: wp.array2d[float],
@@ -239,7 +251,6 @@ def _forward_shim(
     light_pos0: wp.array2d[wp.vec3],
     light_poscom0: wp.array2d[wp.vec3],
     light_targetbodyid: wp.array[int],
-    mapM2M: wp.array[int],
     mat_rgba: wp.array2d[wp.vec4],
     max_ten_J_rownnz: int,
     mesh_face: wp.array[wp.vec3i],
@@ -263,7 +274,6 @@ def _forward_shim(
     mesh_vert: wp.array[wp.vec3],
     mesh_vertadr: wp.array[int],
     mesh_vertnum: wp.array[int],
-    nC: int,
     nJten: int,
     na: int,
     nacttrnbody: int,
@@ -278,6 +288,7 @@ def _forward_shim(
     nflexvert: int,
     ngeom: int,
     ngravcomp: int,
+    nhistory: int,
     njnt: int,
     nlight: int,
     nmaxcondim: int,
@@ -314,1988 +325,6 @@ def _forward_shim(
     qLD_all_updates: wp.array[wp.vec3i],
     qLD_level_offsets: wp.array[int],
     qLD_updates: tuple[wp.array[wp.vec3i], ...],
-    qM_fullm_i: wp.array[int],
-    qM_fullm_j: wp.array[int],
-    qM_mulm_col: wp.array[int],
-    qM_mulm_madr: wp.array[int],
-    qM_mulm_rowadr: wp.array[int],
-    qM_tiles: tuple[mjwp_types.TileSet, ...],
-    qpos0: wp.array2d[float],
-    qpos_spring: wp.array2d[float],
-    rangefinder_sensor_adr: wp.array[int],
-    sensor_acc_adr: wp.array[int],
-    sensor_adr: wp.array[int],
-    sensor_adr_to_contact_adr: wp.array[int],
-    sensor_contact_adr: wp.array[int],
-    sensor_cutoff: wp.array[float],
-    sensor_datatype: wp.array[int],
-    sensor_dim: wp.array[int],
-    sensor_e_kinetic: bool,
-    sensor_e_potential: bool,
-    sensor_intprm: wp.array2d[int],
-    sensor_limitfrc_adr: wp.array[int],
-    sensor_limitpos_adr: wp.array[int],
-    sensor_limitvel_adr: wp.array[int],
-    sensor_objid: wp.array[int],
-    sensor_objtype: wp.array[int],
-    sensor_pos_adr: wp.array[int],
-    sensor_rangefinder_adr: wp.array[int],
-    sensor_rangefinder_bodyid: wp.array[int],
-    sensor_refid: wp.array[int],
-    sensor_reftype: wp.array[int],
-    sensor_rne_postconstraint: bool,
-    sensor_subtree_vel: bool,
-    sensor_tendonactfrc_adr: wp.array[int],
-    sensor_touch_adr: wp.array[int],
-    sensor_type: wp.array[int],
-    sensor_vel_adr: wp.array[int],
-    site_bodyid: wp.array[int],
-    site_pos: wp.array2d[wp.vec3],
-    site_quat: wp.array2d[wp.quat],
-    site_size: wp.array[wp.vec3],
-    site_type: wp.array[int],
-    taxel_sensorid: wp.array[int],
-    taxel_vertadr: wp.array[int],
-    ten_J_colind: wp.array[int],
-    ten_J_rowadr: wp.array[int],
-    ten_J_rownnz: wp.array[int],
-    tendon_actfrclimited: wp.array[bool],
-    tendon_actfrcrange: wp.array2d[wp.vec2],
-    tendon_adr: wp.array[int],
-    tendon_armature: wp.array2d[float],
-    tendon_damping: wp.array2d[float],
-    tendon_dampingpoly: wp.array2d[wp.vec2],
-    tendon_frictionloss: wp.array2d[float],
-    tendon_geom_adr: wp.array[int],
-    tendon_invweight0: wp.array2d[float],
-    tendon_jnt_adr: wp.array[int],
-    tendon_length0: wp.array2d[float],
-    tendon_lengthspring: wp.array2d[wp.vec2],
-    tendon_limited_adr: wp.array[int],
-    tendon_margin: wp.array2d[float],
-    tendon_num: wp.array[int],
-    tendon_range: wp.array2d[wp.vec2],
-    tendon_site_pair_adr: wp.array[int],
-    tendon_solimp_fri: wp.array2d[mjwp_types.vec5],
-    tendon_solimp_lim: wp.array2d[mjwp_types.vec5],
-    tendon_solref_fri: wp.array2d[wp.vec2],
-    tendon_solref_lim: wp.array2d[wp.vec2],
-    tendon_stiffness: wp.array2d[float],
-    tendon_stiffnesspoly: wp.array2d[wp.vec2],
-    wrap_geom_adr: wp.array[int],
-    wrap_jnt_adr: wp.array[int],
-    wrap_objid: wp.array[int],
-    wrap_prm: wp.array[float],
-    wrap_pulley_scale: wp.array[float],
-    wrap_site_pair_adr: wp.array[int],
-    wrap_type: wp.array[int],
-    opt__broadphase: int,
-    opt__broadphase_filter: int,
-    opt__ccd_iterations: int,
-    opt__ccd_tolerance: wp.array[float],
-    opt__cone: int,
-    opt__contact_sensor_maxmatch: int,
-    opt__density: wp.array[float],
-    opt__disableflags: int,
-    opt__enableflags: int,
-    opt__graph_conditional: bool,
-    opt__gravity: wp.array[wp.vec3],
-    opt__impratio_invsqrt: wp.array[float],
-    opt__iterations: int,
-    opt__ls_iterations: int,
-    opt__ls_parallel: bool,
-    opt__ls_parallel_min_step: float,
-    opt__ls_tolerance: wp.array[float],
-    opt__magnetic: wp.array[wp.vec3],
-    opt__run_collision_detection: bool,
-    opt__sdf_initpoints: int,
-    opt__sdf_iterations: int,
-    opt__solver: int,
-    opt__timestep: wp.array[float],
-    opt__tolerance: wp.array[float],
-    opt__viscosity: wp.array[float],
-    opt__wind: wp.array[wp.vec3],
-    stat__meaninertia: wp.array[float],
-    # Data
-    naccdmax: int,
-    naconmax: int,
-    njmax: int,
-    njmax_nnz: int,
-    act: wp.array2d[float],
-    act_dot: wp.array2d[float],
-    actuator_force: wp.array2d[float],
-    actuator_length: wp.array2d[float],
-    actuator_moment: wp.array2d[float],
-    actuator_velocity: wp.array2d[float],
-    cacc: wp.array2d[wp.spatial_vector],
-    cam_xmat: wp.array2d[wp.mat33],
-    cam_xpos: wp.array2d[wp.vec3],
-    cdof: wp.array2d[wp.spatial_vector],
-    cdof_dot: wp.array2d[wp.spatial_vector],
-    cfrc_ext: wp.array2d[wp.spatial_vector],
-    cfrc_int: wp.array2d[wp.spatial_vector],
-    cinert: wp.array2d[mjwp_types.vec10],
-    crb: wp.array2d[mjwp_types.vec10],
-    ctrl: wp.array2d[float],
-    cvel: wp.array2d[wp.spatial_vector],
-    energy: wp.array[wp.vec2],
-    eq_active: wp.array2d[bool],
-    flexedge_J: wp.array2d[float],
-    flexedge_length: wp.array2d[float],
-    flexedge_velocity: wp.array2d[float],
-    flexvert_xpos: wp.array2d[wp.vec3],
-    geom_xmat: wp.array2d[wp.mat33],
-    geom_xpos: wp.array2d[wp.vec3],
-    light_xdir: wp.array2d[wp.vec3],
-    light_xpos: wp.array2d[wp.vec3],
-    mocap_pos: wp.array2d[wp.vec3],
-    mocap_quat: wp.array2d[wp.quat],
-    moment_colind: wp.array2d[int],
-    moment_rowadr: wp.array2d[int],
-    moment_rownnz: wp.array2d[int],
-    nacon: wp.array[int],
-    ncollision: wp.array[int],
-    ne: wp.array[int],
-    nefc: wp.array[int],
-    nf: wp.array[int],
-    nisland: wp.array[int],
-    nl: wp.array[int],
-    qLD: wp.array3d[float],
-    qLDiagInv: wp.array2d[float],
-    qM: wp.array3d[float],
-    qacc: wp.array2d[float],
-    qacc_smooth: wp.array2d[float],
-    qacc_warmstart: wp.array2d[float],
-    qfrc_actuator: wp.array2d[float],
-    qfrc_applied: wp.array2d[float],
-    qfrc_bias: wp.array2d[float],
-    qfrc_constraint: wp.array2d[float],
-    qfrc_damper: wp.array2d[float],
-    qfrc_fluid: wp.array2d[float],
-    qfrc_gravcomp: wp.array2d[float],
-    qfrc_passive: wp.array2d[float],
-    qfrc_smooth: wp.array2d[float],
-    qfrc_spring: wp.array2d[float],
-    qpos: wp.array2d[float],
-    qvel: wp.array2d[float],
-    sensordata: wp.array2d[float],
-    site_xmat: wp.array2d[wp.mat33],
-    site_xpos: wp.array2d[wp.vec3],
-    solver_niter: wp.array[int],
-    subtree_angmom: wp.array2d[wp.vec3],
-    subtree_com: wp.array2d[wp.vec3],
-    subtree_linvel: wp.array2d[wp.vec3],
-    ten_J: wp.array2d[float],
-    ten_length: wp.array2d[float],
-    ten_velocity: wp.array2d[float],
-    ten_wrapadr: wp.array2d[int],
-    ten_wrapnum: wp.array2d[int],
-    time: wp.array[float],
-    tree_island: wp.array2d[int],
-    wrap_obj: wp.array2d[wp.vec2i],
-    wrap_xpos: wp.array2d[wp.spatial_vector],
-    xanchor: wp.array2d[wp.vec3],
-    xaxis: wp.array2d[wp.vec3],
-    xfrc_applied: wp.array2d[wp.spatial_vector],
-    ximat: wp.array2d[wp.mat33],
-    xipos: wp.array2d[wp.vec3],
-    xmat: wp.array2d[wp.mat33],
-    xpos: wp.array2d[wp.vec3],
-    xquat: wp.array2d[wp.quat],
-    contact__dim: wp.array[int],
-    contact__dist: wp.array[float],
-    contact__efc_address: wp.array2d[int],
-    contact__flex: wp.array[wp.vec2i],
-    contact__frame: wp.array[wp.mat33],
-    contact__friction: wp.array[mjwp_types.vec5],
-    contact__geom: wp.array[wp.vec2i],
-    contact__geomcollisionid: wp.array[int],
-    contact__includemargin: wp.array[float],
-    contact__pos: wp.array[wp.vec3],
-    contact__solimp: wp.array[mjwp_types.vec5],
-    contact__solref: wp.array[wp.vec2],
-    contact__solreffriction: wp.array[wp.vec2],
-    contact__type: wp.array[int],
-    contact__vert: wp.array[wp.vec2i],
-    contact__worldid: wp.array[int],
-    efc__D: wp.array2d[float],
-    efc__J: wp.array3d[float],
-    efc__J_colind: wp.array3d[int],
-    efc__J_rowadr: wp.array2d[int],
-    efc__J_rownnz: wp.array2d[int],
-    efc__Jqvel: wp.array2d[float],
-    efc__Ma: wp.array2d[float],
-    efc__aref: wp.array2d[float],
-    efc__force: wp.array2d[float],
-    efc__frictionloss: wp.array2d[float],
-    efc__id: wp.array2d[int],
-    efc__margin: wp.array2d[float],
-    efc__pos: wp.array2d[float],
-    efc__state: wp.array2d[int],
-    efc__type: wp.array2d[int],
-    efc__vel: wp.array2d[float],
-):
-  _m.stat = _s
-  _m.opt = _o
-  _m.callback = _cb
-  _d.efc = _e
-  _d.contact = _c
-  _m.M_rowadr = M_rowadr
-  _m.M_rownnz = M_rownnz
-  _m.actuator_acc0 = actuator_acc0
-  _m.actuator_actadr = actuator_actadr
-  _m.actuator_actearly = actuator_actearly
-  _m.actuator_actlimited = actuator_actlimited
-  _m.actuator_actnum = actuator_actnum
-  _m.actuator_actrange = actuator_actrange
-  _m.actuator_biasprm = actuator_biasprm
-  _m.actuator_biastype = actuator_biastype
-  _m.actuator_cranklength = actuator_cranklength
-  _m.actuator_ctrllimited = actuator_ctrllimited
-  _m.actuator_ctrlrange = actuator_ctrlrange
-  _m.actuator_dynprm = actuator_dynprm
-  _m.actuator_dyntype = actuator_dyntype
-  _m.actuator_forcelimited = actuator_forcelimited
-  _m.actuator_forcerange = actuator_forcerange
-  _m.actuator_gainprm = actuator_gainprm
-  _m.actuator_gaintype = actuator_gaintype
-  _m.actuator_gear = actuator_gear
-  _m.actuator_lengthrange = actuator_lengthrange
-  _m.actuator_trnid = actuator_trnid
-  _m.actuator_trntype = actuator_trntype
-  _m.actuator_trntype_body_adr = actuator_trntype_body_adr
-  _m.block_dim = block_dim
-  _m.body_branch_start = body_branch_start
-  _m.body_branches = body_branches
-  _m.body_dofadr = body_dofadr
-  _m.body_dofnum = body_dofnum
-  _m.body_fluid_ellipsoid = body_fluid_ellipsoid
-  _m.body_geomadr = body_geomadr
-  _m.body_geomnum = body_geomnum
-  _m.body_gravcomp = body_gravcomp
-  _m.body_inertia = body_inertia
-  _m.body_invweight0 = body_invweight0
-  _m.body_ipos = body_ipos
-  _m.body_iquat = body_iquat
-  _m.body_isdofancestor = body_isdofancestor
-  _m.body_jntadr = body_jntadr
-  _m.body_jntnum = body_jntnum
-  _m.body_mass = body_mass
-  _m.body_mocapid = body_mocapid
-  _m.body_parentid = body_parentid
-  _m.body_pos = body_pos
-  _m.body_quat = body_quat
-  _m.body_rootid = body_rootid
-  _m.body_subtreemass = body_subtreemass
-  _m.body_tree = body_tree
-  _m.body_treeid = body_treeid
-  _m.body_weldid = body_weldid
-  _m.cam_bodyid = cam_bodyid
-  _m.cam_fovy = cam_fovy
-  _m.cam_intrinsic = cam_intrinsic
-  _m.cam_mat0 = cam_mat0
-  _m.cam_mode = cam_mode
-  _m.cam_pos = cam_pos
-  _m.cam_pos0 = cam_pos0
-  _m.cam_poscom0 = cam_poscom0
-  _m.cam_quat = cam_quat
-  _m.cam_resolution = cam_resolution
-  _m.cam_sensorsize = cam_sensorsize
-  _m.cam_targetbodyid = cam_targetbodyid
-  _m.dof_Madr = dof_Madr
-  _m.dof_armature = dof_armature
-  _m.dof_bodyid = dof_bodyid
-  _m.dof_damping = dof_damping
-  _m.dof_dampingpoly = dof_dampingpoly
-  _m.dof_frictionloss = dof_frictionloss
-  _m.dof_invweight0 = dof_invweight0
-  _m.dof_jntid = dof_jntid
-  _m.dof_parentid = dof_parentid
-  _m.dof_solimp = dof_solimp
-  _m.dof_solref = dof_solref
-  _m.dof_treeid = dof_treeid
-  _m.dof_tri_col = dof_tri_col
-  _m.dof_tri_row = dof_tri_row
-  _m.eq_connect_adr = eq_connect_adr
-  _m.eq_data = eq_data
-  _m.eq_flex_adr = eq_flex_adr
-  _m.eq_jnt_adr = eq_jnt_adr
-  _m.eq_obj1id = eq_obj1id
-  _m.eq_obj2id = eq_obj2id
-  _m.eq_objtype = eq_objtype
-  _m.eq_solimp = eq_solimp
-  _m.eq_solref = eq_solref
-  _m.eq_ten_adr = eq_ten_adr
-  _m.eq_type = eq_type
-  _m.eq_wld_adr = eq_wld_adr
-  _m.flex_bending = flex_bending
-  _m.flex_bendingadr = flex_bendingadr
-  _m.flex_centered = flex_centered
-  _m.flex_conaffinity = flex_conaffinity
-  _m.flex_condim = flex_condim
-  _m.flex_contype = flex_contype
-  _m.flex_damping = flex_damping
-  _m.flex_dim = flex_dim
-  _m.flex_edge = flex_edge
-  _m.flex_edgeadr = flex_edgeadr
-  _m.flex_edgeflap = flex_edgeflap
-  _m.flex_edgenum = flex_edgenum
-  _m.flex_elem = flex_elem
-  _m.flex_elemadr = flex_elemadr
-  _m.flex_elemdataadr = flex_elemdataadr
-  _m.flex_elemedge = flex_elemedge
-  _m.flex_elemedgeadr = flex_elemedgeadr
-  _m.flex_elemnum = flex_elemnum
-  _m.flex_friction = flex_friction
-  _m.flex_gap = flex_gap
-  _m.flex_margin = flex_margin
-  _m.flex_priority = flex_priority
-  _m.flex_radius = flex_radius
-  _m.flex_shell = flex_shell
-  _m.flex_shelldataadr = flex_shelldataadr
-  _m.flex_shellnum = flex_shellnum
-  _m.flex_solimp = flex_solimp
-  _m.flex_solmix = flex_solmix
-  _m.flex_solref = flex_solref
-  _m.flex_stiffness = flex_stiffness
-  _m.flex_stiffnessadr = flex_stiffnessadr
-  _m.flex_vert = flex_vert
-  _m.flex_vertadr = flex_vertadr
-  _m.flex_vertbodyid = flex_vertbodyid
-  _m.flex_vertflexid = flex_vertflexid
-  _m.flex_vertnum = flex_vertnum
-  _m.flexedge_J_colind = flexedge_J_colind
-  _m.flexedge_J_rowadr = flexedge_J_rowadr
-  _m.flexedge_J_rownnz = flexedge_J_rownnz
-  _m.flexedge_invweight0 = flexedge_invweight0
-  _m.flexedge_length0 = flexedge_length0
-  _m.geom_aabb = geom_aabb
-  _m.geom_bodyid = geom_bodyid
-  _m.geom_conaffinity = geom_conaffinity
-  _m.geom_condim = geom_condim
-  _m.geom_contype = geom_contype
-  _m.geom_dataid = geom_dataid
-  _m.geom_fluid = geom_fluid
-  _m.geom_friction = geom_friction
-  _m.geom_gap = geom_gap
-  _m.geom_group = geom_group
-  _m.geom_margin = geom_margin
-  _m.geom_matid = geom_matid
-  _m.geom_pair_type_count = geom_pair_type_count
-  _m.geom_plugin_index = geom_plugin_index
-  _m.geom_pos = geom_pos
-  _m.geom_priority = geom_priority
-  _m.geom_quat = geom_quat
-  _m.geom_rbound = geom_rbound
-  _m.geom_rgba = geom_rgba
-  _m.geom_size = geom_size
-  _m.geom_solimp = geom_solimp
-  _m.geom_solmix = geom_solmix
-  _m.geom_solref = geom_solref
-  _m.geom_type = geom_type
-  _m.has_fluid = has_fluid
-  _m.has_sdf_geom = has_sdf_geom
-  _m.hfield_adr = hfield_adr
-  _m.hfield_data = hfield_data
-  _m.hfield_ncol = hfield_ncol
-  _m.hfield_nrow = hfield_nrow
-  _m.hfield_size = hfield_size
-  _m.is_sparse = is_sparse
-  _m.jnt_actfrclimited = jnt_actfrclimited
-  _m.jnt_actfrcrange = jnt_actfrcrange
-  _m.jnt_actgravcomp = jnt_actgravcomp
-  _m.jnt_axis = jnt_axis
-  _m.jnt_bodyid = jnt_bodyid
-  _m.jnt_dofadr = jnt_dofadr
-  _m.jnt_limited_ball_adr = jnt_limited_ball_adr
-  _m.jnt_limited_slide_hinge_adr = jnt_limited_slide_hinge_adr
-  _m.jnt_margin = jnt_margin
-  _m.jnt_pos = jnt_pos
-  _m.jnt_qposadr = jnt_qposadr
-  _m.jnt_range = jnt_range
-  _m.jnt_solimp = jnt_solimp
-  _m.jnt_solref = jnt_solref
-  _m.jnt_stiffness = jnt_stiffness
-  _m.jnt_stiffnesspoly = jnt_stiffnesspoly
-  _m.jnt_type = jnt_type
-  _m.light_bodyid = light_bodyid
-  _m.light_dir = light_dir
-  _m.light_dir0 = light_dir0
-  _m.light_mode = light_mode
-  _m.light_pos = light_pos
-  _m.light_pos0 = light_pos0
-  _m.light_poscom0 = light_poscom0
-  _m.light_targetbodyid = light_targetbodyid
-  _m.mapM2M = mapM2M
-  _m.mat_rgba = mat_rgba
-  _m.max_ten_J_rownnz = max_ten_J_rownnz
-  _m.mesh_face = mesh_face
-  _m.mesh_faceadr = mesh_faceadr
-  _m.mesh_graph = mesh_graph
-  _m.mesh_graphadr = mesh_graphadr
-  _m.mesh_normal = mesh_normal
-  _m.mesh_normaladr = mesh_normaladr
-  _m.mesh_normalnum = mesh_normalnum
-  _m.mesh_octadr = mesh_octadr
-  _m.mesh_polyadr = mesh_polyadr
-  _m.mesh_polymap = mesh_polymap
-  _m.mesh_polymapadr = mesh_polymapadr
-  _m.mesh_polymapnum = mesh_polymapnum
-  _m.mesh_polynormal = mesh_polynormal
-  _m.mesh_polynum = mesh_polynum
-  _m.mesh_polyvert = mesh_polyvert
-  _m.mesh_polyvertadr = mesh_polyvertadr
-  _m.mesh_polyvertnum = mesh_polyvertnum
-  _m.mesh_quat = mesh_quat
-  _m.mesh_vert = mesh_vert
-  _m.mesh_vertadr = mesh_vertadr
-  _m.mesh_vertnum = mesh_vertnum
-  _m.nC = nC
-  _m.nJten = nJten
-  _m.na = na
-  _m.nacttrnbody = nacttrnbody
-  _m.nbody = nbody
-  _m.nbranch = nbranch
-  _m.ncam = ncam
-  _m.neq = neq
-  _m.nflex = nflex
-  _m.nflexedge = nflexedge
-  _m.nflexelem = nflexelem
-  _m.nflexshelldata = nflexshelldata
-  _m.nflexvert = nflexvert
-  _m.ngeom = ngeom
-  _m.ngravcomp = ngravcomp
-  _m.njnt = njnt
-  _m.nlight = nlight
-  _m.nmaxcondim = nmaxcondim
-  _m.nmaxmeshdeg = nmaxmeshdeg
-  _m.nmaxpolygon = nmaxpolygon
-  _m.nmaxpyramid = nmaxpyramid
-  _m.nmeshface = nmeshface
-  _m.nrangefinder = nrangefinder
-  _m.nsensorcollision = nsensorcollision
-  _m.nsensorcontact = nsensorcontact
-  _m.nsensortaxel = nsensortaxel
-  _m.nsite = nsite
-  _m.ntendon = ntendon
-  _m.ntree = ntree
-  _m.nu = nu
-  _m.nv = nv
-  _m.nv_pad = nv_pad
-  _m.nwrap = nwrap
-  _m.nxn_geom_pair_filtered = nxn_geom_pair_filtered
-  _m.nxn_pairid = nxn_pairid
-  _m.nxn_pairid_filtered = nxn_pairid_filtered
-  _m.oct_aabb = oct_aabb
-  _m.oct_child = oct_child
-  _m.oct_coeff = oct_coeff
-  _m.opt.broadphase = opt__broadphase
-  _m.opt.broadphase_filter = opt__broadphase_filter
-  _m.opt.ccd_iterations = opt__ccd_iterations
-  _m.opt.ccd_tolerance = opt__ccd_tolerance
-  _m.opt.cone = opt__cone
-  _m.opt.contact_sensor_maxmatch = opt__contact_sensor_maxmatch
-  _m.opt.density = opt__density
-  _m.opt.disableflags = opt__disableflags
-  _m.opt.enableflags = opt__enableflags
-  _m.opt.graph_conditional = opt__graph_conditional
-  _m.opt.gravity = opt__gravity
-  _m.opt.impratio_invsqrt = opt__impratio_invsqrt
-  _m.opt.iterations = opt__iterations
-  _m.opt.ls_iterations = opt__ls_iterations
-  _m.opt.ls_parallel = opt__ls_parallel
-  _m.opt.ls_parallel_min_step = opt__ls_parallel_min_step
-  _m.opt.ls_tolerance = opt__ls_tolerance
-  _m.opt.magnetic = opt__magnetic
-  _m.opt.run_collision_detection = opt__run_collision_detection
-  _m.opt.sdf_initpoints = opt__sdf_initpoints
-  _m.opt.sdf_iterations = opt__sdf_iterations
-  _m.opt.solver = opt__solver
-  _m.opt.timestep = opt__timestep
-  _m.opt.tolerance = opt__tolerance
-  _m.opt.viscosity = opt__viscosity
-  _m.opt.wind = opt__wind
-  _m.pair_dim = pair_dim
-  _m.pair_friction = pair_friction
-  _m.pair_gap = pair_gap
-  _m.pair_margin = pair_margin
-  _m.pair_solimp = pair_solimp
-  _m.pair_solref = pair_solref
-  _m.pair_solreffriction = pair_solreffriction
-  _m.plugin = plugin
-  _m.plugin_attr = plugin_attr
-  _m.qLD_all_updates = qLD_all_updates
-  _m.qLD_level_offsets = qLD_level_offsets
-  _m.qLD_updates = qLD_updates
-  _m.qM_fullm_i = qM_fullm_i
-  _m.qM_fullm_j = qM_fullm_j
-  _m.qM_mulm_col = qM_mulm_col
-  _m.qM_mulm_madr = qM_mulm_madr
-  _m.qM_mulm_rowadr = qM_mulm_rowadr
-  _m.qM_tiles = qM_tiles
-  _m.qpos0 = qpos0
-  _m.qpos_spring = qpos_spring
-  _m.rangefinder_sensor_adr = rangefinder_sensor_adr
-  _m.sensor_acc_adr = sensor_acc_adr
-  _m.sensor_adr = sensor_adr
-  _m.sensor_adr_to_contact_adr = sensor_adr_to_contact_adr
-  _m.sensor_contact_adr = sensor_contact_adr
-  _m.sensor_cutoff = sensor_cutoff
-  _m.sensor_datatype = sensor_datatype
-  _m.sensor_dim = sensor_dim
-  _m.sensor_e_kinetic = sensor_e_kinetic
-  _m.sensor_e_potential = sensor_e_potential
-  _m.sensor_intprm = sensor_intprm
-  _m.sensor_limitfrc_adr = sensor_limitfrc_adr
-  _m.sensor_limitpos_adr = sensor_limitpos_adr
-  _m.sensor_limitvel_adr = sensor_limitvel_adr
-  _m.sensor_objid = sensor_objid
-  _m.sensor_objtype = sensor_objtype
-  _m.sensor_pos_adr = sensor_pos_adr
-  _m.sensor_rangefinder_adr = sensor_rangefinder_adr
-  _m.sensor_rangefinder_bodyid = sensor_rangefinder_bodyid
-  _m.sensor_refid = sensor_refid
-  _m.sensor_reftype = sensor_reftype
-  _m.sensor_rne_postconstraint = sensor_rne_postconstraint
-  _m.sensor_subtree_vel = sensor_subtree_vel
-  _m.sensor_tendonactfrc_adr = sensor_tendonactfrc_adr
-  _m.sensor_touch_adr = sensor_touch_adr
-  _m.sensor_type = sensor_type
-  _m.sensor_vel_adr = sensor_vel_adr
-  _m.site_bodyid = site_bodyid
-  _m.site_pos = site_pos
-  _m.site_quat = site_quat
-  _m.site_size = site_size
-  _m.site_type = site_type
-  _m.stat.meaninertia = stat__meaninertia
-  _m.taxel_sensorid = taxel_sensorid
-  _m.taxel_vertadr = taxel_vertadr
-  _m.ten_J_colind = ten_J_colind
-  _m.ten_J_rowadr = ten_J_rowadr
-  _m.ten_J_rownnz = ten_J_rownnz
-  _m.tendon_actfrclimited = tendon_actfrclimited
-  _m.tendon_actfrcrange = tendon_actfrcrange
-  _m.tendon_adr = tendon_adr
-  _m.tendon_armature = tendon_armature
-  _m.tendon_damping = tendon_damping
-  _m.tendon_dampingpoly = tendon_dampingpoly
-  _m.tendon_frictionloss = tendon_frictionloss
-  _m.tendon_geom_adr = tendon_geom_adr
-  _m.tendon_invweight0 = tendon_invweight0
-  _m.tendon_jnt_adr = tendon_jnt_adr
-  _m.tendon_length0 = tendon_length0
-  _m.tendon_lengthspring = tendon_lengthspring
-  _m.tendon_limited_adr = tendon_limited_adr
-  _m.tendon_margin = tendon_margin
-  _m.tendon_num = tendon_num
-  _m.tendon_range = tendon_range
-  _m.tendon_site_pair_adr = tendon_site_pair_adr
-  _m.tendon_solimp_fri = tendon_solimp_fri
-  _m.tendon_solimp_lim = tendon_solimp_lim
-  _m.tendon_solref_fri = tendon_solref_fri
-  _m.tendon_solref_lim = tendon_solref_lim
-  _m.tendon_stiffness = tendon_stiffness
-  _m.tendon_stiffnesspoly = tendon_stiffnesspoly
-  _m.wrap_geom_adr = wrap_geom_adr
-  _m.wrap_jnt_adr = wrap_jnt_adr
-  _m.wrap_objid = wrap_objid
-  _m.wrap_prm = wrap_prm
-  _m.wrap_pulley_scale = wrap_pulley_scale
-  _m.wrap_site_pair_adr = wrap_site_pair_adr
-  _m.wrap_type = wrap_type
-  _d.act = act
-  _d.act_dot = act_dot
-  _d.actuator_force = actuator_force
-  _d.actuator_length = actuator_length
-  _d.actuator_moment = actuator_moment
-  _d.actuator_velocity = actuator_velocity
-  _d.cacc = cacc
-  _d.cam_xmat = cam_xmat
-  _d.cam_xpos = cam_xpos
-  _d.cdof = cdof
-  _d.cdof_dot = cdof_dot
-  _d.cfrc_ext = cfrc_ext
-  _d.cfrc_int = cfrc_int
-  _d.cinert = cinert
-  _d.contact.dim = contact__dim
-  _d.contact.dist = contact__dist
-  _d.contact.efc_address = contact__efc_address
-  _d.contact.flex = contact__flex
-  _d.contact.frame = contact__frame
-  _d.contact.friction = contact__friction
-  _d.contact.geom = contact__geom
-  _d.contact.geomcollisionid = contact__geomcollisionid
-  _d.contact.includemargin = contact__includemargin
-  _d.contact.pos = contact__pos
-  _d.contact.solimp = contact__solimp
-  _d.contact.solref = contact__solref
-  _d.contact.solreffriction = contact__solreffriction
-  _d.contact.type = contact__type
-  _d.contact.vert = contact__vert
-  _d.contact.worldid = contact__worldid
-  _d.crb = crb
-  _d.ctrl = ctrl
-  _d.cvel = cvel
-  _d.efc.D = efc__D
-  _d.efc.J = efc__J
-  _d.efc.J_colind = efc__J_colind
-  _d.efc.J_rowadr = efc__J_rowadr
-  _d.efc.J_rownnz = efc__J_rownnz
-  _d.efc.Jqvel = efc__Jqvel
-  _d.efc.Ma = efc__Ma
-  _d.efc.aref = efc__aref
-  _d.efc.force = efc__force
-  _d.efc.frictionloss = efc__frictionloss
-  _d.efc.id = efc__id
-  _d.efc.margin = efc__margin
-  _d.efc.pos = efc__pos
-  _d.efc.state = efc__state
-  _d.efc.type = efc__type
-  _d.efc.vel = efc__vel
-  _d.energy = energy
-  _d.eq_active = eq_active
-  _d.flexedge_J = flexedge_J
-  _d.flexedge_length = flexedge_length
-  _d.flexedge_velocity = flexedge_velocity
-  _d.flexvert_xpos = flexvert_xpos
-  _d.geom_xmat = geom_xmat
-  _d.geom_xpos = geom_xpos
-  _d.light_xdir = light_xdir
-  _d.light_xpos = light_xpos
-  _d.mocap_pos = mocap_pos
-  _d.mocap_quat = mocap_quat
-  _d.moment_colind = moment_colind
-  _d.moment_rowadr = moment_rowadr
-  _d.moment_rownnz = moment_rownnz
-  _d.naccdmax = naccdmax
-  _d.nacon = nacon
-  _d.naconmax = naconmax
-  _d.ncollision = ncollision
-  _d.ne = ne
-  _d.nefc = nefc
-  _d.nf = nf
-  _d.nisland = nisland
-  _d.njmax = njmax
-  _d.njmax_nnz = njmax_nnz
-  _d.nl = nl
-  _d.qLD = qLD
-  _d.qLDiagInv = qLDiagInv
-  _d.qM = qM
-  _d.qacc = qacc
-  _d.qacc_smooth = qacc_smooth
-  _d.qacc_warmstart = qacc_warmstart
-  _d.qfrc_actuator = qfrc_actuator
-  _d.qfrc_applied = qfrc_applied
-  _d.qfrc_bias = qfrc_bias
-  _d.qfrc_constraint = qfrc_constraint
-  _d.qfrc_damper = qfrc_damper
-  _d.qfrc_fluid = qfrc_fluid
-  _d.qfrc_gravcomp = qfrc_gravcomp
-  _d.qfrc_passive = qfrc_passive
-  _d.qfrc_smooth = qfrc_smooth
-  _d.qfrc_spring = qfrc_spring
-  _d.qpos = qpos
-  _d.qvel = qvel
-  _d.sensordata = sensordata
-  _d.site_xmat = site_xmat
-  _d.site_xpos = site_xpos
-  _d.solver_niter = solver_niter
-  _d.subtree_angmom = subtree_angmom
-  _d.subtree_com = subtree_com
-  _d.subtree_linvel = subtree_linvel
-  _d.ten_J = ten_J
-  _d.ten_length = ten_length
-  _d.ten_velocity = ten_velocity
-  _d.ten_wrapadr = ten_wrapadr
-  _d.ten_wrapnum = ten_wrapnum
-  _d.time = time
-  _d.tree_island = tree_island
-  _d.wrap_obj = wrap_obj
-  _d.wrap_xpos = wrap_xpos
-  _d.xanchor = xanchor
-  _d.xaxis = xaxis
-  _d.xfrc_applied = xfrc_applied
-  _d.ximat = ximat
-  _d.xipos = xipos
-  _d.xmat = xmat
-  _d.xpos = xpos
-  _d.xquat = xquat
-  _d.nworld = nworld
-  mjwarp.forward(_m, _d)
-
-
-def _forward_jax_impl(m: types.Model, d: types.Data):
-  output_dims = {
-      'act_dot': d.act_dot.shape,
-      'actuator_force': d.actuator_force.shape,
-      'actuator_length': d.actuator_length.shape,
-      'actuator_moment': d._impl.actuator_moment.shape,
-      'actuator_velocity': d._impl.actuator_velocity.shape,
-      'cacc': d._impl.cacc.shape,
-      'cam_xmat': d.cam_xmat.shape,
-      'cam_xpos': d.cam_xpos.shape,
-      'cdof': d.cdof.shape,
-      'cdof_dot': d.cdof_dot.shape,
-      'cfrc_ext': d._impl.cfrc_ext.shape,
-      'cfrc_int': d._impl.cfrc_int.shape,
-      'cinert': d._impl.cinert.shape,
-      'crb': d._impl.crb.shape,
-      'cvel': d.cvel.shape,
-      'energy': d._impl.energy.shape,
-      'flexedge_J': d._impl.flexedge_J.shape,
-      'flexedge_length': d._impl.flexedge_length.shape,
-      'flexedge_velocity': d._impl.flexedge_velocity.shape,
-      'flexvert_xpos': d._impl.flexvert_xpos.shape,
-      'geom_xmat': d.geom_xmat.shape,
-      'geom_xpos': d.geom_xpos.shape,
-      'light_xdir': d._impl.light_xdir.shape,
-      'light_xpos': d._impl.light_xpos.shape,
-      'moment_colind': d._impl.moment_colind.shape,
-      'moment_rowadr': d._impl.moment_rowadr.shape,
-      'moment_rownnz': d._impl.moment_rownnz.shape,
-      'nacon': d._impl.nacon.shape,
-      'ncollision': d._impl.ncollision.shape,
-      'ne': d._impl.ne.shape,
-      'nefc': d._impl.nefc.shape,
-      'nf': d._impl.nf.shape,
-      'nisland': d._impl.nisland.shape,
-      'nl': d._impl.nl.shape,
-      'qLD': d._impl.qLD.shape,
-      'qLDiagInv': d._impl.qLDiagInv.shape,
-      'qM': d._impl.qM.shape,
-      'qacc': d.qacc.shape,
-      'qacc_smooth': d.qacc_smooth.shape,
-      'qfrc_actuator': d.qfrc_actuator.shape,
-      'qfrc_bias': d.qfrc_bias.shape,
-      'qfrc_constraint': d.qfrc_constraint.shape,
-      'qfrc_damper': d._impl.qfrc_damper.shape,
-      'qfrc_fluid': d.qfrc_fluid.shape,
-      'qfrc_gravcomp': d.qfrc_gravcomp.shape,
-      'qfrc_passive': d.qfrc_passive.shape,
-      'qfrc_smooth': d.qfrc_smooth.shape,
-      'qfrc_spring': d._impl.qfrc_spring.shape,
-      'qvel': d.qvel.shape,
-      'sensordata': d.sensordata.shape,
-      'site_xmat': d.site_xmat.shape,
-      'site_xpos': d.site_xpos.shape,
-      'solver_niter': d._impl.solver_niter.shape,
-      'subtree_angmom': d._impl.subtree_angmom.shape,
-      'subtree_com': d.subtree_com.shape,
-      'subtree_linvel': d._impl.subtree_linvel.shape,
-      'ten_J': d._impl.ten_J.shape,
-      'ten_length': d.ten_length.shape,
-      'ten_velocity': d._impl.ten_velocity.shape,
-      'ten_wrapadr': d._impl.ten_wrapadr.shape,
-      'ten_wrapnum': d._impl.ten_wrapnum.shape,
-      'tree_island': d._impl.tree_island.shape,
-      'wrap_obj': d._impl.wrap_obj.shape,
-      'wrap_xpos': d._impl.wrap_xpos.shape,
-      'xanchor': d.xanchor.shape,
-      'xaxis': d.xaxis.shape,
-      'ximat': d.ximat.shape,
-      'xipos': d.xipos.shape,
-      'xmat': d.xmat.shape,
-      'xpos': d.xpos.shape,
-      'xquat': d.xquat.shape,
-      'contact__dim': d._impl.contact__dim.shape,
-      'contact__dist': d._impl.contact__dist.shape,
-      'contact__efc_address': d._impl.contact__efc_address.shape,
-      'contact__flex': d._impl.contact__flex.shape,
-      'contact__frame': d._impl.contact__frame.shape,
-      'contact__friction': d._impl.contact__friction.shape,
-      'contact__geom': d._impl.contact__geom.shape,
-      'contact__geomcollisionid': d._impl.contact__geomcollisionid.shape,
-      'contact__includemargin': d._impl.contact__includemargin.shape,
-      'contact__pos': d._impl.contact__pos.shape,
-      'contact__solimp': d._impl.contact__solimp.shape,
-      'contact__solref': d._impl.contact__solref.shape,
-      'contact__solreffriction': d._impl.contact__solreffriction.shape,
-      'contact__type': d._impl.contact__type.shape,
-      'contact__vert': d._impl.contact__vert.shape,
-      'contact__worldid': d._impl.contact__worldid.shape,
-      'efc__D': d._impl.efc__D.shape,
-      'efc__J': d._impl.efc__J.shape,
-      'efc__J_colind': d._impl.efc__J_colind.shape,
-      'efc__J_rowadr': d._impl.efc__J_rowadr.shape,
-      'efc__J_rownnz': d._impl.efc__J_rownnz.shape,
-      'efc__Jqvel': d._impl.efc__Jqvel.shape,
-      'efc__Ma': d._impl.efc__Ma.shape,
-      'efc__aref': d._impl.efc__aref.shape,
-      'efc__force': d._impl.efc__force.shape,
-      'efc__frictionloss': d._impl.efc__frictionloss.shape,
-      'efc__id': d._impl.efc__id.shape,
-      'efc__margin': d._impl.efc__margin.shape,
-      'efc__pos': d._impl.efc__pos.shape,
-      'efc__state': d._impl.efc__state.shape,
-      'efc__type': d._impl.efc__type.shape,
-      'efc__vel': d._impl.efc__vel.shape,
-  }
-  jf = ffi.jax_callable_variadic_tuple(
-      _forward_shim,
-      num_outputs=103,
-      output_dims=output_dims,
-      vmap_method=None,
-      in_out_argnames=set([
-          'act_dot',
-          'actuator_force',
-          'actuator_length',
-          'actuator_moment',
-          'actuator_velocity',
-          'cacc',
-          'cam_xmat',
-          'cam_xpos',
-          'cdof',
-          'cdof_dot',
-          'cfrc_ext',
-          'cfrc_int',
-          'cinert',
-          'crb',
-          'cvel',
-          'energy',
-          'flexedge_J',
-          'flexedge_length',
-          'flexedge_velocity',
-          'flexvert_xpos',
-          'geom_xmat',
-          'geom_xpos',
-          'light_xdir',
-          'light_xpos',
-          'moment_colind',
-          'moment_rowadr',
-          'moment_rownnz',
-          'nacon',
-          'ncollision',
-          'ne',
-          'nefc',
-          'nf',
-          'nisland',
-          'nl',
-          'qLD',
-          'qLDiagInv',
-          'qM',
-          'qacc',
-          'qacc_smooth',
-          'qfrc_actuator',
-          'qfrc_bias',
-          'qfrc_constraint',
-          'qfrc_damper',
-          'qfrc_fluid',
-          'qfrc_gravcomp',
-          'qfrc_passive',
-          'qfrc_smooth',
-          'qfrc_spring',
-          'qvel',
-          'sensordata',
-          'site_xmat',
-          'site_xpos',
-          'solver_niter',
-          'subtree_angmom',
-          'subtree_com',
-          'subtree_linvel',
-          'ten_J',
-          'ten_length',
-          'ten_velocity',
-          'ten_wrapadr',
-          'ten_wrapnum',
-          'tree_island',
-          'wrap_obj',
-          'wrap_xpos',
-          'xanchor',
-          'xaxis',
-          'ximat',
-          'xipos',
-          'xmat',
-          'xpos',
-          'xquat',
-          'contact__dim',
-          'contact__dist',
-          'contact__efc_address',
-          'contact__flex',
-          'contact__frame',
-          'contact__friction',
-          'contact__geom',
-          'contact__geomcollisionid',
-          'contact__includemargin',
-          'contact__pos',
-          'contact__solimp',
-          'contact__solref',
-          'contact__solreffriction',
-          'contact__type',
-          'contact__vert',
-          'contact__worldid',
-          'efc__D',
-          'efc__J',
-          'efc__J_colind',
-          'efc__J_rowadr',
-          'efc__J_rownnz',
-          'efc__Jqvel',
-          'efc__Ma',
-          'efc__aref',
-          'efc__force',
-          'efc__frictionloss',
-          'efc__id',
-          'efc__margin',
-          'efc__pos',
-          'efc__state',
-          'efc__type',
-          'efc__vel',
-      ]),
-      stage_in_argnames=set([
-          'act',
-          'act_dot',
-          'actuator_acc0',
-          'actuator_actrange',
-          'actuator_biasprm',
-          'actuator_cranklength',
-          'actuator_ctrlrange',
-          'actuator_dynprm',
-          'actuator_force',
-          'actuator_forcerange',
-          'actuator_gainprm',
-          'actuator_gear',
-          'actuator_length',
-          'actuator_lengthrange',
-          'body_gravcomp',
-          'body_inertia',
-          'body_invweight0',
-          'body_ipos',
-          'body_iquat',
-          'body_mass',
-          'body_pos',
-          'body_quat',
-          'body_subtreemass',
-          'cam_fovy',
-          'cam_intrinsic',
-          'cam_mat0',
-          'cam_pos',
-          'cam_pos0',
-          'cam_poscom0',
-          'cam_quat',
-          'cam_xmat',
-          'cam_xpos',
-          'cdof',
-          'cdof_dot',
-          'ctrl',
-          'cvel',
-          'dof_armature',
-          'dof_damping',
-          'dof_dampingpoly',
-          'dof_frictionloss',
-          'dof_invweight0',
-          'dof_solimp',
-          'dof_solref',
-          'eq_active',
-          'eq_data',
-          'eq_solimp',
-          'eq_solref',
-          'geom_aabb',
-          'geom_friction',
-          'geom_gap',
-          'geom_margin',
-          'geom_matid',
-          'geom_pos',
-          'geom_quat',
-          'geom_rbound',
-          'geom_rgba',
-          'geom_size',
-          'geom_solimp',
-          'geom_solmix',
-          'geom_solref',
-          'geom_xmat',
-          'geom_xpos',
-          'hfield_data',
-          'jnt_actfrcrange',
-          'jnt_axis',
-          'jnt_margin',
-          'jnt_pos',
-          'jnt_range',
-          'jnt_solimp',
-          'jnt_solref',
-          'jnt_stiffness',
-          'jnt_stiffnesspoly',
-          'light_dir',
-          'light_dir0',
-          'light_pos',
-          'light_pos0',
-          'light_poscom0',
-          'mat_rgba',
-          'mocap_pos',
-          'mocap_quat',
-          'opt__density',
-          'opt__gravity',
-          'opt__ls_tolerance',
-          'opt__magnetic',
-          'opt__timestep',
-          'opt__tolerance',
-          'opt__viscosity',
-          'opt__wind',
-          'pair_friction',
-          'pair_gap',
-          'pair_margin',
-          'pair_solimp',
-          'pair_solref',
-          'pair_solreffriction',
-          'qacc',
-          'qacc_smooth',
-          'qacc_warmstart',
-          'qfrc_actuator',
-          'qfrc_applied',
-          'qfrc_bias',
-          'qfrc_constraint',
-          'qfrc_fluid',
-          'qfrc_gravcomp',
-          'qfrc_passive',
-          'qfrc_smooth',
-          'qpos',
-          'qpos0',
-          'qpos_spring',
-          'qvel',
-          'sensordata',
-          'site_pos',
-          'site_quat',
-          'site_xmat',
-          'site_xpos',
-          'subtree_com',
-          'ten_length',
-          'tendon_actfrcrange',
-          'tendon_armature',
-          'tendon_damping',
-          'tendon_dampingpoly',
-          'tendon_frictionloss',
-          'tendon_invweight0',
-          'tendon_length0',
-          'tendon_lengthspring',
-          'tendon_margin',
-          'tendon_range',
-          'tendon_solimp_fri',
-          'tendon_solimp_lim',
-          'tendon_solref_fri',
-          'tendon_solref_lim',
-          'tendon_stiffness',
-          'tendon_stiffnesspoly',
-          'time',
-          'xanchor',
-          'xaxis',
-          'xfrc_applied',
-          'ximat',
-          'xipos',
-          'xmat',
-          'xpos',
-          'xquat',
-      ]),
-      stage_out_argnames=set([
-          'act_dot',
-          'actuator_force',
-          'actuator_length',
-          'cam_xmat',
-          'cam_xpos',
-          'cdof',
-          'cdof_dot',
-          'cvel',
-          'geom_xmat',
-          'geom_xpos',
-          'qacc',
-          'qacc_smooth',
-          'qfrc_actuator',
-          'qfrc_bias',
-          'qfrc_constraint',
-          'qfrc_fluid',
-          'qfrc_gravcomp',
-          'qfrc_passive',
-          'qfrc_smooth',
-          'qvel',
-          'sensordata',
-          'site_xmat',
-          'site_xpos',
-          'subtree_com',
-          'ten_length',
-          'xanchor',
-          'xaxis',
-          'ximat',
-          'xipos',
-          'xmat',
-          'xpos',
-          'xquat',
-      ]),
-      graph_mode=m.opt._impl.graph_mode,
-      has_side_effect=False,
-  )
-  out = jf(
-      d.qpos.shape[0],
-      m._impl.M_rowadr,
-      m._impl.M_rownnz,
-      m.actuator_acc0,
-      m.actuator_actadr,
-      m.actuator_actearly,
-      m.actuator_actlimited,
-      m.actuator_actnum,
-      m.actuator_actrange,
-      m.actuator_biasprm,
-      m.actuator_biastype,
-      m.actuator_cranklength,
-      m.actuator_ctrllimited,
-      m.actuator_ctrlrange,
-      m.actuator_dynprm,
-      m.actuator_dyntype,
-      m.actuator_forcelimited,
-      m.actuator_forcerange,
-      m.actuator_gainprm,
-      m.actuator_gaintype,
-      m.actuator_gear,
-      m.actuator_lengthrange,
-      m.actuator_trnid,
-      m.actuator_trntype,
-      m._impl.actuator_trntype_body_adr,
-      m._impl.block_dim,
-      m._impl.body_branch_start,
-      m._impl.body_branches,
-      m.body_dofadr,
-      m.body_dofnum,
-      m._impl.body_fluid_ellipsoid,
-      m.body_geomadr,
-      m.body_geomnum,
-      m.body_gravcomp,
-      m.body_inertia,
-      m.body_invweight0,
-      m.body_ipos,
-      m.body_iquat,
-      m._impl.body_isdofancestor,
-      m.body_jntadr,
-      m.body_jntnum,
-      m.body_mass,
-      m.body_mocapid,
-      m.body_parentid,
-      m.body_pos,
-      m.body_quat,
-      m.body_rootid,
-      m.body_subtreemass,
-      m._impl.body_tree,
-      m.body_treeid,
-      m.body_weldid,
-      m.cam_bodyid,
-      m.cam_fovy,
-      m.cam_intrinsic,
-      m.cam_mat0,
-      m.cam_mode,
-      m.cam_pos,
-      m.cam_pos0,
-      m.cam_poscom0,
-      m.cam_quat,
-      m.cam_resolution,
-      m.cam_sensorsize,
-      m.cam_targetbodyid,
-      m.dof_Madr,
-      m.dof_armature,
-      m.dof_bodyid,
-      m.dof_damping,
-      m.dof_dampingpoly,
-      m.dof_frictionloss,
-      m.dof_invweight0,
-      m.dof_jntid,
-      m.dof_parentid,
-      m.dof_solimp,
-      m.dof_solref,
-      m.dof_treeid,
-      m._impl.dof_tri_col,
-      m._impl.dof_tri_row,
-      m._impl.eq_connect_adr,
-      m.eq_data,
-      m._impl.eq_flex_adr,
-      m._impl.eq_jnt_adr,
-      m.eq_obj1id,
-      m.eq_obj2id,
-      m.eq_objtype,
-      m.eq_solimp,
-      m.eq_solref,
-      m._impl.eq_ten_adr,
-      m.eq_type,
-      m._impl.eq_wld_adr,
-      m._impl.flex_bending,
-      m._impl.flex_bendingadr,
-      m._impl.flex_centered,
-      m._impl.flex_conaffinity,
-      m._impl.flex_condim,
-      m._impl.flex_contype,
-      m._impl.flex_damping,
-      m._impl.flex_dim,
-      m._impl.flex_edge,
-      m._impl.flex_edgeadr,
-      m._impl.flex_edgeflap,
-      m._impl.flex_edgenum,
-      m._impl.flex_elem,
-      m._impl.flex_elemadr,
-      m._impl.flex_elemdataadr,
-      m._impl.flex_elemedge,
-      m._impl.flex_elemedgeadr,
-      m._impl.flex_elemnum,
-      m._impl.flex_friction,
-      m._impl.flex_gap,
-      m._impl.flex_margin,
-      m._impl.flex_priority,
-      m._impl.flex_radius,
-      m._impl.flex_shell,
-      m._impl.flex_shelldataadr,
-      m._impl.flex_shellnum,
-      m._impl.flex_solimp,
-      m._impl.flex_solmix,
-      m._impl.flex_solref,
-      m._impl.flex_stiffness,
-      m._impl.flex_stiffnessadr,
-      m._impl.flex_vert,
-      m.flex_vertadr,
-      m._impl.flex_vertbodyid,
-      m._impl.flex_vertflexid,
-      m.flex_vertnum,
-      m._impl.flexedge_J_colind,
-      m._impl.flexedge_J_rowadr,
-      m._impl.flexedge_J_rownnz,
-      m._impl.flexedge_invweight0,
-      m._impl.flexedge_length0,
-      m.geom_aabb,
-      m.geom_bodyid,
-      m.geom_conaffinity,
-      m.geom_condim,
-      m.geom_contype,
-      jax.numpy.expand_dims(m.geom_dataid, 0),
-      m.geom_fluid,
-      m.geom_friction,
-      m.geom_gap,
-      m.geom_group,
-      m.geom_margin,
-      m.geom_matid,
-      m._impl.geom_pair_type_count,
-      m._impl.geom_plugin_index,
-      m.geom_pos,
-      m.geom_priority,
-      m.geom_quat,
-      m.geom_rbound,
-      m.geom_rgba,
-      m.geom_size,
-      m.geom_solimp,
-      m.geom_solmix,
-      m.geom_solref,
-      m.geom_type,
-      m._impl.has_fluid,
-      m._impl.has_sdf_geom,
-      m.hfield_adr,
-      m.hfield_data,
-      m.hfield_ncol,
-      m.hfield_nrow,
-      m.hfield_size,
-      m._impl.is_sparse,
-      m.jnt_actfrclimited,
-      m.jnt_actfrcrange,
-      m.jnt_actgravcomp,
-      m.jnt_axis,
-      m.jnt_bodyid,
-      m.jnt_dofadr,
-      m._impl.jnt_limited_ball_adr,
-      m._impl.jnt_limited_slide_hinge_adr,
-      m.jnt_margin,
-      m.jnt_pos,
-      m.jnt_qposadr,
-      m.jnt_range,
-      m.jnt_solimp,
-      m.jnt_solref,
-      m.jnt_stiffness,
-      m.jnt_stiffnesspoly,
-      m.jnt_type,
-      m._impl.light_bodyid,
-      m.light_dir,
-      m.light_dir0,
-      m.light_mode,
-      m.light_pos,
-      m.light_pos0,
-      m.light_poscom0,
-      m._impl.light_targetbodyid,
-      m._impl.mapM2M,
-      m.mat_rgba,
-      m._impl.max_ten_J_rownnz,
-      m.mesh_face,
-      m.mesh_faceadr,
-      m.mesh_graph,
-      m.mesh_graphadr,
-      m.mesh_normal,
-      m.mesh_normaladr,
-      m.mesh_normalnum,
-      m.mesh_octadr,
-      m._impl.mesh_polyadr,
-      m._impl.mesh_polymap,
-      m._impl.mesh_polymapadr,
-      m._impl.mesh_polymapnum,
-      m._impl.mesh_polynormal,
-      m._impl.mesh_polynum,
-      m._impl.mesh_polyvert,
-      m._impl.mesh_polyvertadr,
-      m._impl.mesh_polyvertnum,
-      m.mesh_quat,
-      m.mesh_vert,
-      m.mesh_vertadr,
-      m.mesh_vertnum,
-      m.nC,
-      m.nJten,
-      m.na,
-      m._impl.nacttrnbody,
-      m.nbody,
-      m._impl.nbranch,
-      m.ncam,
-      m.neq,
-      m.nflex,
-      m._impl.nflexedge,
-      m._impl.nflexelem,
-      m._impl.nflexshelldata,
-      m._impl.nflexvert,
-      m.ngeom,
-      m.ngravcomp,
-      m.njnt,
-      m.nlight,
-      m._impl.nmaxcondim,
-      m._impl.nmaxmeshdeg,
-      m._impl.nmaxpolygon,
-      m._impl.nmaxpyramid,
-      m.nmeshface,
-      m._impl.nrangefinder,
-      m._impl.nsensorcollision,
-      m._impl.nsensorcontact,
-      m._impl.nsensortaxel,
-      m.nsite,
-      m.ntendon,
-      m._impl.ntree,
-      m.nu,
-      m.nv,
-      m._impl.nv_pad,
-      m.nwrap,
-      m._impl.nxn_geom_pair_filtered,
-      m._impl.nxn_pairid,
-      m._impl.nxn_pairid_filtered,
-      m._impl.oct_aabb,
-      m._impl.oct_child,
-      m._impl.oct_coeff,
-      m.pair_dim,
-      m.pair_friction,
-      m.pair_gap,
-      m.pair_margin,
-      m.pair_solimp,
-      m.pair_solref,
-      m.pair_solreffriction,
-      m._impl.plugin,
-      m._impl.plugin_attr,
-      m._impl.qLD_all_updates,
-      m._impl.qLD_level_offsets,
-      m._impl.qLD_updates,
-      m._impl.qM_fullm_i,
-      m._impl.qM_fullm_j,
-      m._impl.qM_mulm_col,
-      m._impl.qM_mulm_madr,
-      m._impl.qM_mulm_rowadr,
-      m._impl.qM_tiles,
-      m.qpos0,
-      m.qpos_spring,
-      m._impl.rangefinder_sensor_adr,
-      m._impl.sensor_acc_adr,
-      m.sensor_adr,
-      m._impl.sensor_adr_to_contact_adr,
-      m._impl.sensor_contact_adr,
-      m.sensor_cutoff,
-      m.sensor_datatype,
-      m.sensor_dim,
-      m._impl.sensor_e_kinetic,
-      m._impl.sensor_e_potential,
-      m.sensor_intprm,
-      m._impl.sensor_limitfrc_adr,
-      m._impl.sensor_limitpos_adr,
-      m._impl.sensor_limitvel_adr,
-      m.sensor_objid,
-      m.sensor_objtype,
-      m._impl.sensor_pos_adr,
-      m._impl.sensor_rangefinder_adr,
-      m._impl.sensor_rangefinder_bodyid,
-      m.sensor_refid,
-      m.sensor_reftype,
-      m._impl.sensor_rne_postconstraint,
-      m._impl.sensor_subtree_vel,
-      m._impl.sensor_tendonactfrc_adr,
-      m._impl.sensor_touch_adr,
-      m.sensor_type,
-      m._impl.sensor_vel_adr,
-      m.site_bodyid,
-      m.site_pos,
-      m.site_quat,
-      m.site_size,
-      m.site_type,
-      m._impl.taxel_sensorid,
-      m._impl.taxel_vertadr,
-      m._impl.ten_J_colind,
-      m._impl.ten_J_rowadr,
-      m._impl.ten_J_rownnz,
-      m.tendon_actfrclimited,
-      m.tendon_actfrcrange,
-      m.tendon_adr,
-      m.tendon_armature,
-      m.tendon_damping,
-      m.tendon_dampingpoly,
-      m.tendon_frictionloss,
-      m._impl.tendon_geom_adr,
-      m.tendon_invweight0,
-      m._impl.tendon_jnt_adr,
-      m.tendon_length0,
-      m.tendon_lengthspring,
-      m._impl.tendon_limited_adr,
-      m.tendon_margin,
-      m.tendon_num,
-      m.tendon_range,
-      m._impl.tendon_site_pair_adr,
-      m.tendon_solimp_fri,
-      m.tendon_solimp_lim,
-      m.tendon_solref_fri,
-      m.tendon_solref_lim,
-      m.tendon_stiffness,
-      m.tendon_stiffnesspoly,
-      m._impl.wrap_geom_adr,
-      m._impl.wrap_jnt_adr,
-      m.wrap_objid,
-      m.wrap_prm,
-      m._impl.wrap_pulley_scale,
-      m._impl.wrap_site_pair_adr,
-      m.wrap_type,
-      m.opt._impl.broadphase,
-      m.opt._impl.broadphase_filter,
-      m.opt._impl.ccd_iterations,
-      m.opt._impl.ccd_tolerance,
-      m.opt.cone,
-      m.opt._impl.contact_sensor_maxmatch,
-      m.opt.density,
-      m.opt.disableflags,
-      m.opt.enableflags,
-      m.opt._impl.graph_conditional,
-      m.opt.gravity,
-      m.opt._impl.impratio_invsqrt,
-      m.opt.iterations,
-      m.opt.ls_iterations,
-      m.opt._impl.ls_parallel,
-      m.opt._impl.ls_parallel_min_step,
-      m.opt.ls_tolerance,
-      m.opt.magnetic,
-      m.opt._impl.run_collision_detection,
-      m.opt._impl.sdf_initpoints,
-      m.opt._impl.sdf_iterations,
-      m.opt.solver,
-      m.opt.timestep,
-      m.opt.tolerance,
-      m.opt.viscosity,
-      m.opt.wind,
-      m.stat.meaninertia,
-      d._impl.naccdmax,
-      d._impl.naconmax,
-      d._impl.njmax,
-      d._impl.njmax_nnz,
-      d.act,
-      d.act_dot,
-      d.actuator_force,
-      d.actuator_length,
-      d._impl.actuator_moment,
-      d._impl.actuator_velocity,
-      d._impl.cacc,
-      d.cam_xmat,
-      d.cam_xpos,
-      d.cdof,
-      d.cdof_dot,
-      d._impl.cfrc_ext,
-      d._impl.cfrc_int,
-      d._impl.cinert,
-      d._impl.crb,
-      d.ctrl,
-      d.cvel,
-      d._impl.energy,
-      d.eq_active,
-      d._impl.flexedge_J,
-      d._impl.flexedge_length,
-      d._impl.flexedge_velocity,
-      d._impl.flexvert_xpos,
-      d.geom_xmat,
-      d.geom_xpos,
-      d._impl.light_xdir,
-      d._impl.light_xpos,
-      d.mocap_pos,
-      d.mocap_quat,
-      d._impl.moment_colind,
-      d._impl.moment_rowadr,
-      d._impl.moment_rownnz,
-      d._impl.nacon,
-      d._impl.ncollision,
-      d._impl.ne,
-      d._impl.nefc,
-      d._impl.nf,
-      d._impl.nisland,
-      d._impl.nl,
-      d._impl.qLD,
-      d._impl.qLDiagInv,
-      d._impl.qM,
-      d.qacc,
-      d.qacc_smooth,
-      d.qacc_warmstart,
-      d.qfrc_actuator,
-      d.qfrc_applied,
-      d.qfrc_bias,
-      d.qfrc_constraint,
-      d._impl.qfrc_damper,
-      d.qfrc_fluid,
-      d.qfrc_gravcomp,
-      d.qfrc_passive,
-      d.qfrc_smooth,
-      d._impl.qfrc_spring,
-      d.qpos,
-      d.qvel,
-      d.sensordata,
-      d.site_xmat,
-      d.site_xpos,
-      d._impl.solver_niter,
-      d._impl.subtree_angmom,
-      d.subtree_com,
-      d._impl.subtree_linvel,
-      d._impl.ten_J,
-      d.ten_length,
-      d._impl.ten_velocity,
-      d._impl.ten_wrapadr,
-      d._impl.ten_wrapnum,
-      d.time,
-      d._impl.tree_island,
-      d._impl.wrap_obj,
-      d._impl.wrap_xpos,
-      d.xanchor,
-      d.xaxis,
-      d.xfrc_applied,
-      d.ximat,
-      d.xipos,
-      d.xmat,
-      d.xpos,
-      d.xquat,
-      d._impl.contact__dim,
-      d._impl.contact__dist,
-      d._impl.contact__efc_address,
-      d._impl.contact__flex,
-      d._impl.contact__frame,
-      d._impl.contact__friction,
-      d._impl.contact__geom,
-      d._impl.contact__geomcollisionid,
-      d._impl.contact__includemargin,
-      d._impl.contact__pos,
-      d._impl.contact__solimp,
-      d._impl.contact__solref,
-      d._impl.contact__solreffriction,
-      d._impl.contact__type,
-      d._impl.contact__vert,
-      d._impl.contact__worldid,
-      d._impl.efc__D,
-      d._impl.efc__J,
-      d._impl.efc__J_colind,
-      d._impl.efc__J_rowadr,
-      d._impl.efc__J_rownnz,
-      d._impl.efc__Jqvel,
-      d._impl.efc__Ma,
-      d._impl.efc__aref,
-      d._impl.efc__force,
-      d._impl.efc__frictionloss,
-      d._impl.efc__id,
-      d._impl.efc__margin,
-      d._impl.efc__pos,
-      d._impl.efc__state,
-      d._impl.efc__type,
-      d._impl.efc__vel,
-  )
-  d = d.tree_replace({
-      'act_dot': out[0],
-      'actuator_force': out[1],
-      'actuator_length': out[2],
-      '_impl.actuator_moment': out[3],
-      '_impl.actuator_velocity': out[4],
-      '_impl.cacc': out[5],
-      'cam_xmat': out[6],
-      'cam_xpos': out[7],
-      'cdof': out[8],
-      'cdof_dot': out[9],
-      '_impl.cfrc_ext': out[10],
-      '_impl.cfrc_int': out[11],
-      '_impl.cinert': out[12],
-      '_impl.crb': out[13],
-      'cvel': out[14],
-      '_impl.energy': out[15],
-      '_impl.flexedge_J': out[16],
-      '_impl.flexedge_length': out[17],
-      '_impl.flexedge_velocity': out[18],
-      '_impl.flexvert_xpos': out[19],
-      'geom_xmat': out[20],
-      'geom_xpos': out[21],
-      '_impl.light_xdir': out[22],
-      '_impl.light_xpos': out[23],
-      '_impl.moment_colind': out[24],
-      '_impl.moment_rowadr': out[25],
-      '_impl.moment_rownnz': out[26],
-      '_impl.nacon': out[27],
-      '_impl.ncollision': out[28],
-      '_impl.ne': out[29],
-      '_impl.nefc': out[30],
-      '_impl.nf': out[31],
-      '_impl.nisland': out[32],
-      '_impl.nl': out[33],
-      '_impl.qLD': out[34],
-      '_impl.qLDiagInv': out[35],
-      '_impl.qM': out[36],
-      'qacc': out[37],
-      'qacc_smooth': out[38],
-      'qfrc_actuator': out[39],
-      'qfrc_bias': out[40],
-      'qfrc_constraint': out[41],
-      '_impl.qfrc_damper': out[42],
-      'qfrc_fluid': out[43],
-      'qfrc_gravcomp': out[44],
-      'qfrc_passive': out[45],
-      'qfrc_smooth': out[46],
-      '_impl.qfrc_spring': out[47],
-      'qvel': out[48],
-      'sensordata': out[49],
-      'site_xmat': out[50],
-      'site_xpos': out[51],
-      '_impl.solver_niter': out[52],
-      '_impl.subtree_angmom': out[53],
-      'subtree_com': out[54],
-      '_impl.subtree_linvel': out[55],
-      '_impl.ten_J': out[56],
-      'ten_length': out[57],
-      '_impl.ten_velocity': out[58],
-      '_impl.ten_wrapadr': out[59],
-      '_impl.ten_wrapnum': out[60],
-      '_impl.tree_island': out[61],
-      '_impl.wrap_obj': out[62],
-      '_impl.wrap_xpos': out[63],
-      'xanchor': out[64],
-      'xaxis': out[65],
-      'ximat': out[66],
-      'xipos': out[67],
-      'xmat': out[68],
-      'xpos': out[69],
-      'xquat': out[70],
-      '_impl.contact__dim': out[71],
-      '_impl.contact__dist': out[72],
-      '_impl.contact__efc_address': out[73],
-      '_impl.contact__flex': out[74],
-      '_impl.contact__frame': out[75],
-      '_impl.contact__friction': out[76],
-      '_impl.contact__geom': out[77],
-      '_impl.contact__geomcollisionid': out[78],
-      '_impl.contact__includemargin': out[79],
-      '_impl.contact__pos': out[80],
-      '_impl.contact__solimp': out[81],
-      '_impl.contact__solref': out[82],
-      '_impl.contact__solreffriction': out[83],
-      '_impl.contact__type': out[84],
-      '_impl.contact__vert': out[85],
-      '_impl.contact__worldid': out[86],
-      '_impl.efc__D': out[87],
-      '_impl.efc__J': out[88],
-      '_impl.efc__J_colind': out[89],
-      '_impl.efc__J_rowadr': out[90],
-      '_impl.efc__J_rownnz': out[91],
-      '_impl.efc__Jqvel': out[92],
-      '_impl.efc__Ma': out[93],
-      '_impl.efc__aref': out[94],
-      '_impl.efc__force': out[95],
-      '_impl.efc__frictionloss': out[96],
-      '_impl.efc__id': out[97],
-      '_impl.efc__margin': out[98],
-      '_impl.efc__pos': out[99],
-      '_impl.efc__state': out[100],
-      '_impl.efc__type': out[101],
-      '_impl.efc__vel': out[102],
-  })
-  return d
-
-
-@jax.custom_batching.custom_vmap
-@ffi.marshal_jax_warp_callable
-def forward(m: types.Model, d: types.Data):
-  return _forward_jax_impl(m, d)
-
-
-@forward.def_vmap
-@ffi.marshal_custom_vmap
-def forward_vmap(unused_axis_size, is_batched, m: types.Model, d: types.Data):
-  d = forward(m, d)
-  return d, is_batched[1]
-
-
-@ffi.format_args_for_warp
-def _step_shim(
-    # Model
-    nworld: int,
-    M_rowadr: wp.array[int],
-    M_rownnz: wp.array[int],
-    actuator_acc0: wp.array2d[float],
-    actuator_actadr: wp.array[int],
-    actuator_actearly: wp.array[bool],
-    actuator_actlimited: wp.array[bool],
-    actuator_actnum: wp.array[int],
-    actuator_actrange: wp.array2d[wp.vec2],
-    actuator_biasprm: wp.array2d[mjwp_types.vec10f],
-    actuator_biastype: wp.array[int],
-    actuator_cranklength: wp.array2d[float],
-    actuator_ctrllimited: wp.array[bool],
-    actuator_ctrlrange: wp.array2d[wp.vec2],
-    actuator_dynprm: wp.array2d[mjwp_types.vec10f],
-    actuator_dyntype: wp.array[int],
-    actuator_forcelimited: wp.array[bool],
-    actuator_forcerange: wp.array2d[wp.vec2],
-    actuator_gainprm: wp.array2d[mjwp_types.vec10f],
-    actuator_gaintype: wp.array[int],
-    actuator_gear: wp.array2d[wp.spatial_vector],
-    actuator_lengthrange: wp.array2d[wp.vec2],
-    actuator_trnid: wp.array[wp.vec2i],
-    actuator_trntype: wp.array[int],
-    actuator_trntype_body_adr: wp.array[int],
-    block_dim: mjwp_types.BlockDim,
-    body_branch_start: wp.array[int],
-    body_branches: wp.array[int],
-    body_dofadr: wp.array[int],
-    body_dofnum: wp.array[int],
-    body_fluid_ellipsoid: wp.array[bool],
-    body_geomadr: wp.array[int],
-    body_geomnum: wp.array[int],
-    body_gravcomp: wp.array2d[float],
-    body_inertia: wp.array2d[wp.vec3],
-    body_invweight0: wp.array2d[wp.vec2],
-    body_ipos: wp.array2d[wp.vec3],
-    body_iquat: wp.array2d[wp.quat],
-    body_isdofancestor: wp.array2d[int],
-    body_jntadr: wp.array[int],
-    body_jntnum: wp.array[int],
-    body_mass: wp.array2d[float],
-    body_mocapid: wp.array[int],
-    body_parentid: wp.array[int],
-    body_pos: wp.array2d[wp.vec3],
-    body_quat: wp.array2d[wp.quat],
-    body_rootid: wp.array[int],
-    body_subtreemass: wp.array2d[float],
-    body_tree: tuple[wp.array[int], ...],
-    body_treeid: wp.array[int],
-    body_weldid: wp.array[int],
-    cam_bodyid: wp.array[int],
-    cam_fovy: wp.array2d[float],
-    cam_intrinsic: wp.array2d[wp.vec4],
-    cam_mat0: wp.array2d[wp.mat33],
-    cam_mode: wp.array[int],
-    cam_pos: wp.array2d[wp.vec3],
-    cam_pos0: wp.array2d[wp.vec3],
-    cam_poscom0: wp.array2d[wp.vec3],
-    cam_quat: wp.array2d[wp.quat],
-    cam_resolution: wp.array[wp.vec2i],
-    cam_sensorsize: wp.array[wp.vec2],
-    cam_targetbodyid: wp.array[int],
-    dof_Madr: wp.array[int],
-    dof_armature: wp.array2d[float],
-    dof_bodyid: wp.array[int],
-    dof_damping: wp.array2d[float],
-    dof_dampingpoly: wp.array2d[wp.vec2],
-    dof_frictionloss: wp.array2d[float],
-    dof_invweight0: wp.array2d[float],
-    dof_jntid: wp.array[int],
-    dof_parentid: wp.array[int],
-    dof_solimp: wp.array2d[mjwp_types.vec5],
-    dof_solref: wp.array2d[wp.vec2],
-    dof_treeid: wp.array[int],
-    dof_tri_col: wp.array[int],
-    dof_tri_row: wp.array[int],
-    eq_connect_adr: wp.array[int],
-    eq_data: wp.array2d[mjwp_types.vec11],
-    eq_flex_adr: wp.array[int],
-    eq_jnt_adr: wp.array[int],
-    eq_obj1id: wp.array[int],
-    eq_obj2id: wp.array[int],
-    eq_objtype: wp.array[int],
-    eq_solimp: wp.array2d[mjwp_types.vec5],
-    eq_solref: wp.array2d[wp.vec2],
-    eq_ten_adr: wp.array[int],
-    eq_type: wp.array[int],
-    eq_wld_adr: wp.array[int],
-    flex_bending: wp.array[float],
-    flex_bendingadr: wp.array[int],
-    flex_centered: wp.array[bool],
-    flex_conaffinity: wp.array[int],
-    flex_condim: wp.array[int],
-    flex_contype: wp.array[int],
-    flex_damping: wp.array[float],
-    flex_dim: wp.array[int],
-    flex_edge: wp.array[wp.vec2i],
-    flex_edgeadr: wp.array[int],
-    flex_edgeflap: wp.array[wp.vec2i],
-    flex_edgenum: wp.array[int],
-    flex_elem: wp.array[int],
-    flex_elemadr: wp.array[int],
-    flex_elemdataadr: wp.array[int],
-    flex_elemedge: wp.array[int],
-    flex_elemedgeadr: wp.array[int],
-    flex_elemnum: wp.array[int],
-    flex_friction: wp.array[wp.vec3],
-    flex_gap: wp.array[float],
-    flex_margin: wp.array[float],
-    flex_priority: wp.array[int],
-    flex_radius: wp.array[float],
-    flex_shell: wp.array[int],
-    flex_shelldataadr: wp.array[int],
-    flex_shellnum: wp.array[int],
-    flex_solimp: wp.array[mjwp_types.vec5],
-    flex_solmix: wp.array[float],
-    flex_solref: wp.array[wp.vec2],
-    flex_stiffness: wp.array[float],
-    flex_stiffnessadr: wp.array[int],
-    flex_vert: wp.array[wp.vec3],
-    flex_vertadr: wp.array[int],
-    flex_vertbodyid: wp.array[int],
-    flex_vertflexid: wp.array[int],
-    flex_vertnum: wp.array[int],
-    flexedge_J_colind: wp.array[int],
-    flexedge_J_rowadr: wp.array[int],
-    flexedge_J_rownnz: wp.array[int],
-    flexedge_invweight0: wp.array[float],
-    flexedge_length0: wp.array[float],
-    geom_aabb: wp.array3d[wp.vec3],
-    geom_bodyid: wp.array[int],
-    geom_conaffinity: wp.array[int],
-    geom_condim: wp.array[int],
-    geom_contype: wp.array[int],
-    geom_dataid: wp.array2d[int],
-    geom_fluid: wp.array2d[float],
-    geom_friction: wp.array2d[wp.vec3],
-    geom_gap: wp.array2d[float],
-    geom_group: wp.array[int],
-    geom_margin: wp.array2d[float],
-    geom_matid: wp.array2d[int],
-    geom_pair_type_count: tuple[int, ...],
-    geom_plugin_index: wp.array[int],
-    geom_pos: wp.array2d[wp.vec3],
-    geom_priority: wp.array[int],
-    geom_quat: wp.array2d[wp.quat],
-    geom_rbound: wp.array2d[float],
-    geom_rgba: wp.array2d[wp.vec4],
-    geom_size: wp.array2d[wp.vec3],
-    geom_solimp: wp.array2d[mjwp_types.vec5],
-    geom_solmix: wp.array2d[float],
-    geom_solref: wp.array2d[wp.vec2],
-    geom_type: wp.array[int],
-    has_fluid: bool,
-    has_sdf_geom: bool,
-    hfield_adr: wp.array[int],
-    hfield_data: wp.array[float],
-    hfield_ncol: wp.array[int],
-    hfield_nrow: wp.array[int],
-    hfield_size: wp.array[wp.vec4],
-    is_sparse: bool,
-    jnt_actfrclimited: wp.array[bool],
-    jnt_actfrcrange: wp.array2d[wp.vec2],
-    jnt_actgravcomp: wp.array[int],
-    jnt_axis: wp.array2d[wp.vec3],
-    jnt_bodyid: wp.array[int],
-    jnt_dofadr: wp.array[int],
-    jnt_limited_ball_adr: wp.array[int],
-    jnt_limited_slide_hinge_adr: wp.array[int],
-    jnt_margin: wp.array2d[float],
-    jnt_pos: wp.array2d[wp.vec3],
-    jnt_qposadr: wp.array[int],
-    jnt_range: wp.array2d[wp.vec2],
-    jnt_solimp: wp.array2d[mjwp_types.vec5],
-    jnt_solref: wp.array2d[wp.vec2],
-    jnt_stiffness: wp.array2d[float],
-    jnt_stiffnesspoly: wp.array2d[wp.vec2],
-    jnt_type: wp.array[int],
-    light_bodyid: wp.array[int],
-    light_dir: wp.array2d[wp.vec3],
-    light_dir0: wp.array2d[wp.vec3],
-    light_mode: wp.array[int],
-    light_pos: wp.array2d[wp.vec3],
-    light_pos0: wp.array2d[wp.vec3],
-    light_poscom0: wp.array2d[wp.vec3],
-    light_targetbodyid: wp.array[int],
-    mapM2M: wp.array[int],
-    mat_rgba: wp.array2d[wp.vec4],
-    max_ten_J_rownnz: int,
-    mesh_face: wp.array[wp.vec3i],
-    mesh_faceadr: wp.array[int],
-    mesh_graph: wp.array[int],
-    mesh_graphadr: wp.array[int],
-    mesh_normal: wp.array[wp.vec3],
-    mesh_normaladr: wp.array[int],
-    mesh_normalnum: wp.array[int],
-    mesh_octadr: wp.array[int],
-    mesh_polyadr: wp.array[int],
-    mesh_polymap: wp.array[int],
-    mesh_polymapadr: wp.array[int],
-    mesh_polymapnum: wp.array[int],
-    mesh_polynormal: wp.array[wp.vec3],
-    mesh_polynum: wp.array[int],
-    mesh_polyvert: wp.array[int],
-    mesh_polyvertadr: wp.array[int],
-    mesh_polyvertnum: wp.array[int],
-    mesh_quat: wp.array[wp.quat],
-    mesh_vert: wp.array[wp.vec3],
-    mesh_vertadr: wp.array[int],
-    mesh_vertnum: wp.array[int],
-    nC: int,
-    nJten: int,
-    nM: int,
-    na: int,
-    nacttrnbody: int,
-    nbody: int,
-    nbranch: int,
-    ncam: int,
-    neq: int,
-    nflex: int,
-    nflexedge: int,
-    nflexelem: int,
-    nflexshelldata: int,
-    nflexvert: int,
-    ngeom: int,
-    ngravcomp: int,
-    njnt: int,
-    nlight: int,
-    nmaxcondim: int,
-    nmaxmeshdeg: int,
-    nmaxpolygon: int,
-    nmaxpyramid: int,
-    nmeshface: int,
-    nrangefinder: int,
-    nsensorcollision: int,
-    nsensorcontact: int,
-    nsensortaxel: int,
-    nsite: int,
-    ntendon: int,
-    ntree: int,
-    nu: int,
-    nv: int,
-    nv_pad: int,
-    nwrap: int,
-    nxn_geom_pair_filtered: wp.array[wp.vec2i],
-    nxn_pairid: wp.array[wp.vec2i],
-    nxn_pairid_filtered: wp.array[wp.vec2i],
-    oct_aabb: wp.array2d[wp.vec3],
-    oct_child: wp.array[mjwp_types.vec8i],
-    oct_coeff: wp.array[mjwp_types.vec8],
-    pair_dim: wp.array[int],
-    pair_friction: wp.array2d[mjwp_types.vec5],
-    pair_gap: wp.array2d[float],
-    pair_margin: wp.array2d[float],
-    pair_solimp: wp.array2d[mjwp_types.vec5],
-    pair_solref: wp.array2d[wp.vec2],
-    pair_solreffriction: wp.array2d[wp.vec2],
-    plugin: wp.array[int],
-    plugin_attr: wp.array[mjwp_types.vec_pluginattr],
-    qLD_all_updates: wp.array[wp.vec3i],
-    qLD_level_offsets: wp.array[int],
-    qLD_updates: tuple[wp.array[wp.vec3i], ...],
-    qM_fullm_i: wp.array[int],
-    qM_fullm_j: wp.array[int],
-    qM_mulm_col: wp.array[int],
-    qM_mulm_madr: wp.array[int],
-    qM_mulm_rowadr: wp.array[int],
-    qM_tiles: tuple[mjwp_types.TileSet, ...],
     qpos0: wp.array2d[float],
     qpos_spring: wp.array2d[float],
     rangefinder_sensor_adr: wp.array[int],
@@ -2305,9 +334,13 @@ def _step_shim(
     sensor_contact_adr: wp.array[int],
     sensor_cutoff: wp.array[float],
     sensor_datatype: wp.array[int],
+    sensor_delay: wp.array[float],
     sensor_dim: wp.array[int],
     sensor_e_kinetic: bool,
     sensor_e_potential: bool,
+    sensor_history: wp.array[wp.vec2i],
+    sensor_historyadr: wp.array[int],
+    sensor_interval: wp.array[wp.vec2],
     sensor_intprm: wp.array2d[int],
     sensor_limitfrc_adr: wp.array[int],
     sensor_limitpos_adr: wp.array[int],
@@ -2398,6 +431,7 @@ def _step_shim(
     naconmax: int,
     njmax: int,
     njmax_nnz: int,
+    M: wp.array3d[float],
     act: wp.array2d[float],
     act_dot: wp.array2d[float],
     actuator_force: wp.array2d[float],
@@ -2415,6 +449,9 @@ def _step_shim(
     crb: wp.array2d[mjwp_types.vec10],
     ctrl: wp.array2d[float],
     cvel: wp.array2d[wp.spatial_vector],
+    dof_island: wp.array2d[int],
+    dof_islandid: wp.array2d[int],
+    efc_islandid: wp.array2d[int],
     energy: wp.array[wp.vec2],
     eq_active: wp.array2d[bool],
     flexedge_J: wp.array2d[float],
@@ -2423,8 +460,23 @@ def _step_shim(
     flexvert_xpos: wp.array2d[wp.vec3],
     geom_xmat: wp.array2d[wp.mat33],
     geom_xpos: wp.array2d[wp.vec3],
+    history: wp.array2d[float],
+    iqacc: wp.array2d[float],
+    iqacc_smooth: wp.array2d[float],
+    iqfrc_constraint: wp.array2d[float],
+    iqfrc_smooth: wp.array2d[float],
+    island_dofadr: wp.array2d[int],
+    island_efcadr: wp.array2d[int],
+    island_ne: wp.array2d[int],
+    island_nefc: wp.array2d[int],
+    island_nf: wp.array2d[int],
+    island_nv: wp.array2d[int],
     light_xdir: wp.array2d[wp.vec3],
     light_xpos: wp.array2d[wp.vec3],
+    map_dof2idof: wp.array2d[int],
+    map_efc2iefc: wp.array2d[int],
+    map_idof2dof: wp.array2d[int],
+    map_iefc2efc: wp.array2d[int],
     mocap_pos: wp.array2d[wp.vec3],
     mocap_quat: wp.array2d[wp.quat],
     moment_colind: wp.array2d[int],
@@ -2435,11 +487,11 @@ def _step_shim(
     ne: wp.array[int],
     nefc: wp.array[int],
     nf: wp.array[int],
+    nidof: wp.array[int],
     nisland: wp.array[int],
     nl: wp.array[int],
     qLD: wp.array3d[float],
     qLDiagInv: wp.array2d[float],
-    qM: wp.array3d[float],
     qacc: wp.array2d[float],
     qacc_smooth: wp.array2d[float],
     qacc_warmstart: wp.array2d[float],
@@ -2505,7 +557,19 @@ def _step_shim(
     efc__aref: wp.array2d[float],
     efc__force: wp.array2d[float],
     efc__frictionloss: wp.array2d[float],
+    efc__iD: wp.array2d[float],
+    efc__iJ: wp.array3d[float],
+    efc__iJ_colind: wp.array3d[int],
+    efc__iJ_rowadr: wp.array2d[int],
+    efc__iJ_rownnz: wp.array2d[int],
+    efc__iaref: wp.array2d[float],
     efc__id: wp.array2d[int],
+    efc__iforce: wp.array2d[float],
+    efc__ifrictionloss: wp.array2d[float],
+    efc__iid: wp.array2d[int],
+    efc__island: wp.array2d[int],
+    efc__istate: wp.array2d[int],
+    efc__itype: wp.array2d[int],
     efc__margin: wp.array2d[float],
     efc__pos: wp.array2d[float],
     efc__state: wp.array2d[int],
@@ -2517,8 +581,18 @@ def _step_shim(
   _m.callback = _cb
   _d.efc = _e
   _d.contact = _c
+  _m.M_elemid = M_elemid
+  _m.M_fullm_i = M_fullm_i
+  _m.M_fullm_j = M_fullm_j
+  _m.M_fullm_upper_elemid = M_fullm_upper_elemid
+  _m.M_fullm_upper_i = M_fullm_upper_i
+  _m.M_fullm_upper_j = M_fullm_upper_j
+  _m.M_mulm_col = M_mulm_col
+  _m.M_mulm_madr = M_mulm_madr
+  _m.M_mulm_rowadr = M_mulm_rowadr
   _m.M_rowadr = M_rowadr
   _m.M_rownnz = M_rownnz
+  _m.M_tiles = M_tiles
   _m.actuator_acc0 = actuator_acc0
   _m.actuator_actadr = actuator_actadr
   _m.actuator_actearly = actuator_actearly
@@ -2530,6 +604,7 @@ def _step_shim(
   _m.actuator_cranklength = actuator_cranklength
   _m.actuator_ctrllimited = actuator_ctrllimited
   _m.actuator_ctrlrange = actuator_ctrlrange
+  _m.actuator_delay = actuator_delay
   _m.actuator_dynprm = actuator_dynprm
   _m.actuator_dyntype = actuator_dyntype
   _m.actuator_forcelimited = actuator_forcelimited
@@ -2537,6 +612,8 @@ def _step_shim(
   _m.actuator_gainprm = actuator_gainprm
   _m.actuator_gaintype = actuator_gaintype
   _m.actuator_gear = actuator_gear
+  _m.actuator_history = actuator_history
+  _m.actuator_historyadr = actuator_historyadr
   _m.actuator_lengthrange = actuator_lengthrange
   _m.actuator_trnid = actuator_trnid
   _m.actuator_trntype = actuator_trntype
@@ -2579,7 +656,6 @@ def _step_shim(
   _m.cam_resolution = cam_resolution
   _m.cam_sensorsize = cam_sensorsize
   _m.cam_targetbodyid = cam_targetbodyid
-  _m.dof_Madr = dof_Madr
   _m.dof_armature = dof_armature
   _m.dof_bodyid = dof_bodyid
   _m.dof_damping = dof_damping
@@ -2703,7 +779,6 @@ def _step_shim(
   _m.light_pos0 = light_pos0
   _m.light_poscom0 = light_poscom0
   _m.light_targetbodyid = light_targetbodyid
-  _m.mapM2M = mapM2M
   _m.mat_rgba = mat_rgba
   _m.max_ten_J_rownnz = max_ten_J_rownnz
   _m.mesh_face = mesh_face
@@ -2727,9 +802,7 @@ def _step_shim(
   _m.mesh_vert = mesh_vert
   _m.mesh_vertadr = mesh_vertadr
   _m.mesh_vertnum = mesh_vertnum
-  _m.nC = nC
   _m.nJten = nJten
-  _m.nM = nM
   _m.na = na
   _m.nacttrnbody = nacttrnbody
   _m.nbody = nbody
@@ -2743,6 +816,7 @@ def _step_shim(
   _m.nflexvert = nflexvert
   _m.ngeom = ngeom
   _m.ngravcomp = ngravcomp
+  _m.nhistory = nhistory
   _m.njnt = njnt
   _m.nlight = nlight
   _m.nmaxcondim = nmaxcondim
@@ -2806,12 +880,6 @@ def _step_shim(
   _m.qLD_all_updates = qLD_all_updates
   _m.qLD_level_offsets = qLD_level_offsets
   _m.qLD_updates = qLD_updates
-  _m.qM_fullm_i = qM_fullm_i
-  _m.qM_fullm_j = qM_fullm_j
-  _m.qM_mulm_col = qM_mulm_col
-  _m.qM_mulm_madr = qM_mulm_madr
-  _m.qM_mulm_rowadr = qM_mulm_rowadr
-  _m.qM_tiles = qM_tiles
   _m.qpos0 = qpos0
   _m.qpos_spring = qpos_spring
   _m.rangefinder_sensor_adr = rangefinder_sensor_adr
@@ -2821,9 +889,13 @@ def _step_shim(
   _m.sensor_contact_adr = sensor_contact_adr
   _m.sensor_cutoff = sensor_cutoff
   _m.sensor_datatype = sensor_datatype
+  _m.sensor_delay = sensor_delay
   _m.sensor_dim = sensor_dim
   _m.sensor_e_kinetic = sensor_e_kinetic
   _m.sensor_e_potential = sensor_e_potential
+  _m.sensor_history = sensor_history
+  _m.sensor_historyadr = sensor_historyadr
+  _m.sensor_interval = sensor_interval
   _m.sensor_intprm = sensor_intprm
   _m.sensor_limitfrc_adr = sensor_limitfrc_adr
   _m.sensor_limitpos_adr = sensor_limitpos_adr
@@ -2882,6 +954,7 @@ def _step_shim(
   _m.wrap_pulley_scale = wrap_pulley_scale
   _m.wrap_site_pair_adr = wrap_site_pair_adr
   _m.wrap_type = wrap_type
+  _d.M = M
   _d.act = act
   _d.act_dot = act_dot
   _d.actuator_force = actuator_force
@@ -2915,6 +988,8 @@ def _step_shim(
   _d.crb = crb
   _d.ctrl = ctrl
   _d.cvel = cvel
+  _d.dof_island = dof_island
+  _d.dof_islandid = dof_islandid
   _d.efc.D = efc__D
   _d.efc.J = efc__J
   _d.efc.J_colind = efc__J_colind
@@ -2925,12 +1000,25 @@ def _step_shim(
   _d.efc.aref = efc__aref
   _d.efc.force = efc__force
   _d.efc.frictionloss = efc__frictionloss
+  _d.efc.iD = efc__iD
+  _d.efc.iJ = efc__iJ
+  _d.efc.iJ_colind = efc__iJ_colind
+  _d.efc.iJ_rowadr = efc__iJ_rowadr
+  _d.efc.iJ_rownnz = efc__iJ_rownnz
+  _d.efc.iaref = efc__iaref
   _d.efc.id = efc__id
+  _d.efc.iforce = efc__iforce
+  _d.efc.ifrictionloss = efc__ifrictionloss
+  _d.efc.iid = efc__iid
+  _d.efc.island = efc__island
+  _d.efc.istate = efc__istate
+  _d.efc.itype = efc__itype
   _d.efc.margin = efc__margin
   _d.efc.pos = efc__pos
   _d.efc.state = efc__state
   _d.efc.type = efc__type
   _d.efc.vel = efc__vel
+  _d.efc_islandid = efc_islandid
   _d.energy = energy
   _d.eq_active = eq_active
   _d.flexedge_J = flexedge_J
@@ -2939,8 +1027,23 @@ def _step_shim(
   _d.flexvert_xpos = flexvert_xpos
   _d.geom_xmat = geom_xmat
   _d.geom_xpos = geom_xpos
+  _d.history = history
+  _d.iqacc = iqacc
+  _d.iqacc_smooth = iqacc_smooth
+  _d.iqfrc_constraint = iqfrc_constraint
+  _d.iqfrc_smooth = iqfrc_smooth
+  _d.island_dofadr = island_dofadr
+  _d.island_efcadr = island_efcadr
+  _d.island_ne = island_ne
+  _d.island_nefc = island_nefc
+  _d.island_nf = island_nf
+  _d.island_nv = island_nv
   _d.light_xdir = light_xdir
   _d.light_xpos = light_xpos
+  _d.map_dof2idof = map_dof2idof
+  _d.map_efc2iefc = map_efc2iefc
+  _d.map_idof2dof = map_idof2dof
+  _d.map_iefc2efc = map_iefc2efc
   _d.mocap_pos = mocap_pos
   _d.mocap_quat = mocap_quat
   _d.moment_colind = moment_colind
@@ -2953,13 +1056,13 @@ def _step_shim(
   _d.ne = ne
   _d.nefc = nefc
   _d.nf = nf
+  _d.nidof = nidof
   _d.nisland = nisland
   _d.njmax = njmax
   _d.njmax_nnz = njmax_nnz
   _d.nl = nl
   _d.qLD = qLD
   _d.qLDiagInv = qLDiagInv
-  _d.qM = qM
   _d.qacc = qacc
   _d.qacc_smooth = qacc_smooth
   _d.qacc_warmstart = qacc_warmstart
@@ -3000,12 +1103,12 @@ def _step_shim(
   _d.xpos = xpos
   _d.xquat = xquat
   _d.nworld = nworld
-  mjwarp.step(_m, _d)
+  mjwarp.forward(_m, _d)
 
 
-def _step_jax_impl(m: types.Model, d: types.Data):
+def _forward_jax_impl(m: types.Model, d: types.Data):
   output_dims = {
-      'act': d.act.shape,
+      'M': d._impl.M.shape,
       'act_dot': d.act_dot.shape,
       'actuator_force': d.actuator_force.shape,
       'actuator_length': d.actuator_length.shape,
@@ -3021,6 +1124,9 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       'cinert': d._impl.cinert.shape,
       'crb': d._impl.crb.shape,
       'cvel': d.cvel.shape,
+      'dof_island': d._impl.dof_island.shape,
+      'dof_islandid': d._impl.dof_islandid.shape,
+      'efc_islandid': d._impl.efc_islandid.shape,
       'energy': d._impl.energy.shape,
       'flexedge_J': d._impl.flexedge_J.shape,
       'flexedge_length': d._impl.flexedge_length.shape,
@@ -3028,8 +1134,23 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       'flexvert_xpos': d._impl.flexvert_xpos.shape,
       'geom_xmat': d.geom_xmat.shape,
       'geom_xpos': d.geom_xpos.shape,
+      'history': d.history.shape,
+      'iqacc': d._impl.iqacc.shape,
+      'iqacc_smooth': d._impl.iqacc_smooth.shape,
+      'iqfrc_constraint': d._impl.iqfrc_constraint.shape,
+      'iqfrc_smooth': d._impl.iqfrc_smooth.shape,
+      'island_dofadr': d._impl.island_dofadr.shape,
+      'island_efcadr': d._impl.island_efcadr.shape,
+      'island_ne': d._impl.island_ne.shape,
+      'island_nefc': d._impl.island_nefc.shape,
+      'island_nf': d._impl.island_nf.shape,
+      'island_nv': d._impl.island_nv.shape,
       'light_xdir': d._impl.light_xdir.shape,
       'light_xpos': d._impl.light_xpos.shape,
+      'map_dof2idof': d._impl.map_dof2idof.shape,
+      'map_efc2iefc': d._impl.map_efc2iefc.shape,
+      'map_idof2dof': d._impl.map_idof2dof.shape,
+      'map_iefc2efc': d._impl.map_iefc2efc.shape,
       'moment_colind': d._impl.moment_colind.shape,
       'moment_rowadr': d._impl.moment_rowadr.shape,
       'moment_rownnz': d._impl.moment_rownnz.shape,
@@ -3038,14 +1159,13 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       'ne': d._impl.ne.shape,
       'nefc': d._impl.nefc.shape,
       'nf': d._impl.nf.shape,
+      'nidof': d._impl.nidof.shape,
       'nisland': d._impl.nisland.shape,
       'nl': d._impl.nl.shape,
       'qLD': d._impl.qLD.shape,
       'qLDiagInv': d._impl.qLDiagInv.shape,
-      'qM': d._impl.qM.shape,
       'qacc': d.qacc.shape,
       'qacc_smooth': d.qacc_smooth.shape,
-      'qacc_warmstart': d.qacc_warmstart.shape,
       'qfrc_actuator': d.qfrc_actuator.shape,
       'qfrc_bias': d.qfrc_bias.shape,
       'qfrc_constraint': d.qfrc_constraint.shape,
@@ -3055,7 +1175,6 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       'qfrc_passive': d.qfrc_passive.shape,
       'qfrc_smooth': d.qfrc_smooth.shape,
       'qfrc_spring': d._impl.qfrc_spring.shape,
-      'qpos': d.qpos.shape,
       'qvel': d.qvel.shape,
       'sensordata': d.sensordata.shape,
       'site_xmat': d.site_xmat.shape,
@@ -3069,7 +1188,6 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       'ten_velocity': d._impl.ten_velocity.shape,
       'ten_wrapadr': d._impl.ten_wrapadr.shape,
       'ten_wrapnum': d._impl.ten_wrapnum.shape,
-      'time': d.time.shape,
       'tree_island': d._impl.tree_island.shape,
       'wrap_obj': d._impl.wrap_obj.shape,
       'wrap_xpos': d._impl.wrap_xpos.shape,
@@ -3106,7 +1224,19 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       'efc__aref': d._impl.efc__aref.shape,
       'efc__force': d._impl.efc__force.shape,
       'efc__frictionloss': d._impl.efc__frictionloss.shape,
+      'efc__iD': d._impl.efc__iD.shape,
+      'efc__iJ': d._impl.efc__iJ.shape,
+      'efc__iJ_colind': d._impl.efc__iJ_colind.shape,
+      'efc__iJ_rowadr': d._impl.efc__iJ_rowadr.shape,
+      'efc__iJ_rownnz': d._impl.efc__iJ_rownnz.shape,
+      'efc__iaref': d._impl.efc__iaref.shape,
       'efc__id': d._impl.efc__id.shape,
+      'efc__iforce': d._impl.efc__iforce.shape,
+      'efc__ifrictionloss': d._impl.efc__ifrictionloss.shape,
+      'efc__iid': d._impl.efc__iid.shape,
+      'efc__island': d._impl.efc__island.shape,
+      'efc__istate': d._impl.efc__istate.shape,
+      'efc__itype': d._impl.efc__itype.shape,
       'efc__margin': d._impl.efc__margin.shape,
       'efc__pos': d._impl.efc__pos.shape,
       'efc__state': d._impl.efc__state.shape,
@@ -3114,12 +1244,12 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       'efc__vel': d._impl.efc__vel.shape,
   }
   jf = ffi.jax_callable_variadic_tuple(
-      _step_shim,
-      num_outputs=107,
+      _forward_shim,
+      num_outputs=134,
       output_dims=output_dims,
       vmap_method=None,
       in_out_argnames=set([
-          'act',
+          'M',
           'act_dot',
           'actuator_force',
           'actuator_length',
@@ -3135,6 +1265,9 @@ def _step_jax_impl(m: types.Model, d: types.Data):
           'cinert',
           'crb',
           'cvel',
+          'dof_island',
+          'dof_islandid',
+          'efc_islandid',
           'energy',
           'flexedge_J',
           'flexedge_length',
@@ -3142,8 +1275,23 @@ def _step_jax_impl(m: types.Model, d: types.Data):
           'flexvert_xpos',
           'geom_xmat',
           'geom_xpos',
+          'history',
+          'iqacc',
+          'iqacc_smooth',
+          'iqfrc_constraint',
+          'iqfrc_smooth',
+          'island_dofadr',
+          'island_efcadr',
+          'island_ne',
+          'island_nefc',
+          'island_nf',
+          'island_nv',
           'light_xdir',
           'light_xpos',
+          'map_dof2idof',
+          'map_efc2iefc',
+          'map_idof2dof',
+          'map_iefc2efc',
           'moment_colind',
           'moment_rowadr',
           'moment_rownnz',
@@ -3152,14 +1300,13 @@ def _step_jax_impl(m: types.Model, d: types.Data):
           'ne',
           'nefc',
           'nf',
+          'nidof',
           'nisland',
           'nl',
           'qLD',
           'qLDiagInv',
-          'qM',
           'qacc',
           'qacc_smooth',
-          'qacc_warmstart',
           'qfrc_actuator',
           'qfrc_bias',
           'qfrc_constraint',
@@ -3169,7 +1316,6 @@ def _step_jax_impl(m: types.Model, d: types.Data):
           'qfrc_passive',
           'qfrc_smooth',
           'qfrc_spring',
-          'qpos',
           'qvel',
           'sensordata',
           'site_xmat',
@@ -3183,7 +1329,6 @@ def _step_jax_impl(m: types.Model, d: types.Data):
           'ten_velocity',
           'ten_wrapadr',
           'ten_wrapnum',
-          'time',
           'tree_island',
           'wrap_obj',
           'wrap_xpos',
@@ -3220,7 +1365,19 @@ def _step_jax_impl(m: types.Model, d: types.Data):
           'efc__aref',
           'efc__force',
           'efc__frictionloss',
+          'efc__iD',
+          'efc__iJ',
+          'efc__iJ_colind',
+          'efc__iJ_rowadr',
+          'efc__iJ_rownnz',
+          'efc__iaref',
           'efc__id',
+          'efc__iforce',
+          'efc__ifrictionloss',
+          'efc__iid',
+          'efc__island',
+          'efc__istate',
+          'efc__itype',
           'efc__margin',
           'efc__pos',
           'efc__state',
@@ -3291,6 +1448,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
           'geom_xmat',
           'geom_xpos',
           'hfield_data',
+          'history',
           'jnt_actfrcrange',
           'jnt_axis',
           'jnt_margin',
@@ -3371,7 +1529,6 @@ def _step_jax_impl(m: types.Model, d: types.Data):
           'xquat',
       ]),
       stage_out_argnames=set([
-          'act',
           'act_dot',
           'actuator_force',
           'actuator_length',
@@ -3382,9 +1539,9 @@ def _step_jax_impl(m: types.Model, d: types.Data):
           'cvel',
           'geom_xmat',
           'geom_xpos',
+          'history',
           'qacc',
           'qacc_smooth',
-          'qacc_warmstart',
           'qfrc_actuator',
           'qfrc_bias',
           'qfrc_constraint',
@@ -3392,14 +1549,12 @@ def _step_jax_impl(m: types.Model, d: types.Data):
           'qfrc_gravcomp',
           'qfrc_passive',
           'qfrc_smooth',
-          'qpos',
           'qvel',
           'sensordata',
           'site_xmat',
           'site_xpos',
           'subtree_com',
           'ten_length',
-          'time',
           'xanchor',
           'xaxis',
           'ximat',
@@ -3413,8 +1568,18 @@ def _step_jax_impl(m: types.Model, d: types.Data):
   )
   out = jf(
       d.qpos.shape[0],
+      m._impl.M_elemid,
+      m._impl.M_fullm_i,
+      m._impl.M_fullm_j,
+      m._impl.M_fullm_upper_elemid,
+      m._impl.M_fullm_upper_i,
+      m._impl.M_fullm_upper_j,
+      m._impl.M_mulm_col,
+      m._impl.M_mulm_madr,
+      m._impl.M_mulm_rowadr,
       m._impl.M_rowadr,
       m._impl.M_rownnz,
+      m._impl.M_tiles,
       m.actuator_acc0,
       m.actuator_actadr,
       m.actuator_actearly,
@@ -3426,6 +1591,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m.actuator_cranklength,
       m.actuator_ctrllimited,
       m.actuator_ctrlrange,
+      m._impl.actuator_delay,
       m.actuator_dynprm,
       m.actuator_dyntype,
       m.actuator_forcelimited,
@@ -3433,6 +1599,8 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m.actuator_gainprm,
       m.actuator_gaintype,
       m.actuator_gear,
+      m._impl.actuator_history,
+      m._impl.actuator_historyadr,
       m.actuator_lengthrange,
       m.actuator_trnid,
       m.actuator_trntype,
@@ -3475,7 +1643,6 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m.cam_resolution,
       m.cam_sensorsize,
       m.cam_targetbodyid,
-      m.dof_Madr,
       m.dof_armature,
       m.dof_bodyid,
       m.dof_damping,
@@ -3599,7 +1766,6 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m.light_pos0,
       m.light_poscom0,
       m._impl.light_targetbodyid,
-      m._impl.mapM2M,
       m.mat_rgba,
       m._impl.max_ten_J_rownnz,
       m.mesh_face,
@@ -3623,9 +1789,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m.mesh_vert,
       m.mesh_vertadr,
       m.mesh_vertnum,
-      m.nC,
       m.nJten,
-      m.nM,
       m.na,
       m._impl.nacttrnbody,
       m.nbody,
@@ -3639,6 +1803,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m._impl.nflexvert,
       m.ngeom,
       m.ngravcomp,
+      m.nhistory,
       m.njnt,
       m.nlight,
       m._impl.nmaxcondim,
@@ -3675,12 +1840,6 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m._impl.qLD_all_updates,
       m._impl.qLD_level_offsets,
       m._impl.qLD_updates,
-      m._impl.qM_fullm_i,
-      m._impl.qM_fullm_j,
-      m._impl.qM_mulm_col,
-      m._impl.qM_mulm_madr,
-      m._impl.qM_mulm_rowadr,
-      m._impl.qM_tiles,
       m.qpos0,
       m.qpos_spring,
       m._impl.rangefinder_sensor_adr,
@@ -3690,9 +1849,13 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m._impl.sensor_contact_adr,
       m.sensor_cutoff,
       m.sensor_datatype,
+      m._impl.sensor_delay,
       m.sensor_dim,
       m._impl.sensor_e_kinetic,
       m._impl.sensor_e_potential,
+      m._impl.sensor_history,
+      m._impl.sensor_historyadr,
+      m._impl.sensor_interval,
       m.sensor_intprm,
       m._impl.sensor_limitfrc_adr,
       m._impl.sensor_limitpos_adr,
@@ -3782,6 +1945,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       d._impl.naconmax,
       d._impl.njmax,
       d._impl.njmax_nnz,
+      d._impl.M,
       d.act,
       d.act_dot,
       d.actuator_force,
@@ -3799,6 +1963,9 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       d._impl.crb,
       d.ctrl,
       d.cvel,
+      d._impl.dof_island,
+      d._impl.dof_islandid,
+      d._impl.efc_islandid,
       d._impl.energy,
       d.eq_active,
       d._impl.flexedge_J,
@@ -3807,8 +1974,23 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       d._impl.flexvert_xpos,
       d.geom_xmat,
       d.geom_xpos,
+      d.history,
+      d._impl.iqacc,
+      d._impl.iqacc_smooth,
+      d._impl.iqfrc_constraint,
+      d._impl.iqfrc_smooth,
+      d._impl.island_dofadr,
+      d._impl.island_efcadr,
+      d._impl.island_ne,
+      d._impl.island_nefc,
+      d._impl.island_nf,
+      d._impl.island_nv,
       d._impl.light_xdir,
       d._impl.light_xpos,
+      d._impl.map_dof2idof,
+      d._impl.map_efc2iefc,
+      d._impl.map_idof2dof,
+      d._impl.map_iefc2efc,
       d.mocap_pos,
       d.mocap_quat,
       d._impl.moment_colind,
@@ -3819,11 +2001,11 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       d._impl.ne,
       d._impl.nefc,
       d._impl.nf,
+      d._impl.nidof,
       d._impl.nisland,
       d._impl.nl,
       d._impl.qLD,
       d._impl.qLDiagInv,
-      d._impl.qM,
       d.qacc,
       d.qacc_smooth,
       d.qacc_warmstart,
@@ -3889,7 +2071,19 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       d._impl.efc__aref,
       d._impl.efc__force,
       d._impl.efc__frictionloss,
+      d._impl.efc__iD,
+      d._impl.efc__iJ,
+      d._impl.efc__iJ_colind,
+      d._impl.efc__iJ_rowadr,
+      d._impl.efc__iJ_rownnz,
+      d._impl.efc__iaref,
       d._impl.efc__id,
+      d._impl.efc__iforce,
+      d._impl.efc__ifrictionloss,
+      d._impl.efc__iid,
+      d._impl.efc__island,
+      d._impl.efc__istate,
+      d._impl.efc__itype,
       d._impl.efc__margin,
       d._impl.efc__pos,
       d._impl.efc__state,
@@ -3897,7 +2091,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       d._impl.efc__vel,
   )
   d = d.tree_replace({
-      'act': out[0],
+      '_impl.M': out[0],
       'act_dot': out[1],
       'actuator_force': out[2],
       'actuator_length': out[3],
@@ -3913,97 +2107,2366 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       '_impl.cinert': out[13],
       '_impl.crb': out[14],
       'cvel': out[15],
-      '_impl.energy': out[16],
-      '_impl.flexedge_J': out[17],
-      '_impl.flexedge_length': out[18],
-      '_impl.flexedge_velocity': out[19],
-      '_impl.flexvert_xpos': out[20],
-      'geom_xmat': out[21],
-      'geom_xpos': out[22],
-      '_impl.light_xdir': out[23],
-      '_impl.light_xpos': out[24],
-      '_impl.moment_colind': out[25],
-      '_impl.moment_rowadr': out[26],
-      '_impl.moment_rownnz': out[27],
-      '_impl.nacon': out[28],
-      '_impl.ncollision': out[29],
-      '_impl.ne': out[30],
-      '_impl.nefc': out[31],
-      '_impl.nf': out[32],
-      '_impl.nisland': out[33],
-      '_impl.nl': out[34],
-      '_impl.qLD': out[35],
-      '_impl.qLDiagInv': out[36],
-      '_impl.qM': out[37],
-      'qacc': out[38],
-      'qacc_smooth': out[39],
-      'qacc_warmstart': out[40],
-      'qfrc_actuator': out[41],
-      'qfrc_bias': out[42],
-      'qfrc_constraint': out[43],
-      '_impl.qfrc_damper': out[44],
-      'qfrc_fluid': out[45],
-      'qfrc_gravcomp': out[46],
-      'qfrc_passive': out[47],
-      'qfrc_smooth': out[48],
-      '_impl.qfrc_spring': out[49],
-      'qpos': out[50],
-      'qvel': out[51],
-      'sensordata': out[52],
-      'site_xmat': out[53],
-      'site_xpos': out[54],
-      '_impl.solver_niter': out[55],
-      '_impl.subtree_angmom': out[56],
-      'subtree_com': out[57],
-      '_impl.subtree_linvel': out[58],
-      '_impl.ten_J': out[59],
-      'ten_length': out[60],
-      '_impl.ten_velocity': out[61],
-      '_impl.ten_wrapadr': out[62],
-      '_impl.ten_wrapnum': out[63],
-      'time': out[64],
-      '_impl.tree_island': out[65],
-      '_impl.wrap_obj': out[66],
-      '_impl.wrap_xpos': out[67],
-      'xanchor': out[68],
-      'xaxis': out[69],
-      'ximat': out[70],
-      'xipos': out[71],
-      'xmat': out[72],
-      'xpos': out[73],
-      'xquat': out[74],
-      '_impl.contact__dim': out[75],
-      '_impl.contact__dist': out[76],
-      '_impl.contact__efc_address': out[77],
-      '_impl.contact__flex': out[78],
-      '_impl.contact__frame': out[79],
-      '_impl.contact__friction': out[80],
-      '_impl.contact__geom': out[81],
-      '_impl.contact__geomcollisionid': out[82],
-      '_impl.contact__includemargin': out[83],
-      '_impl.contact__pos': out[84],
-      '_impl.contact__solimp': out[85],
-      '_impl.contact__solref': out[86],
-      '_impl.contact__solreffriction': out[87],
-      '_impl.contact__type': out[88],
-      '_impl.contact__vert': out[89],
-      '_impl.contact__worldid': out[90],
-      '_impl.efc__D': out[91],
-      '_impl.efc__J': out[92],
-      '_impl.efc__J_colind': out[93],
-      '_impl.efc__J_rowadr': out[94],
-      '_impl.efc__J_rownnz': out[95],
-      '_impl.efc__Jqvel': out[96],
-      '_impl.efc__Ma': out[97],
-      '_impl.efc__aref': out[98],
-      '_impl.efc__force': out[99],
-      '_impl.efc__frictionloss': out[100],
-      '_impl.efc__id': out[101],
-      '_impl.efc__margin': out[102],
-      '_impl.efc__pos': out[103],
-      '_impl.efc__state': out[104],
-      '_impl.efc__type': out[105],
-      '_impl.efc__vel': out[106],
+      '_impl.dof_island': out[16],
+      '_impl.dof_islandid': out[17],
+      '_impl.efc_islandid': out[18],
+      '_impl.energy': out[19],
+      '_impl.flexedge_J': out[20],
+      '_impl.flexedge_length': out[21],
+      '_impl.flexedge_velocity': out[22],
+      '_impl.flexvert_xpos': out[23],
+      'geom_xmat': out[24],
+      'geom_xpos': out[25],
+      'history': out[26],
+      '_impl.iqacc': out[27],
+      '_impl.iqacc_smooth': out[28],
+      '_impl.iqfrc_constraint': out[29],
+      '_impl.iqfrc_smooth': out[30],
+      '_impl.island_dofadr': out[31],
+      '_impl.island_efcadr': out[32],
+      '_impl.island_ne': out[33],
+      '_impl.island_nefc': out[34],
+      '_impl.island_nf': out[35],
+      '_impl.island_nv': out[36],
+      '_impl.light_xdir': out[37],
+      '_impl.light_xpos': out[38],
+      '_impl.map_dof2idof': out[39],
+      '_impl.map_efc2iefc': out[40],
+      '_impl.map_idof2dof': out[41],
+      '_impl.map_iefc2efc': out[42],
+      '_impl.moment_colind': out[43],
+      '_impl.moment_rowadr': out[44],
+      '_impl.moment_rownnz': out[45],
+      '_impl.nacon': out[46],
+      '_impl.ncollision': out[47],
+      '_impl.ne': out[48],
+      '_impl.nefc': out[49],
+      '_impl.nf': out[50],
+      '_impl.nidof': out[51],
+      '_impl.nisland': out[52],
+      '_impl.nl': out[53],
+      '_impl.qLD': out[54],
+      '_impl.qLDiagInv': out[55],
+      'qacc': out[56],
+      'qacc_smooth': out[57],
+      'qfrc_actuator': out[58],
+      'qfrc_bias': out[59],
+      'qfrc_constraint': out[60],
+      '_impl.qfrc_damper': out[61],
+      'qfrc_fluid': out[62],
+      'qfrc_gravcomp': out[63],
+      'qfrc_passive': out[64],
+      'qfrc_smooth': out[65],
+      '_impl.qfrc_spring': out[66],
+      'qvel': out[67],
+      'sensordata': out[68],
+      'site_xmat': out[69],
+      'site_xpos': out[70],
+      '_impl.solver_niter': out[71],
+      '_impl.subtree_angmom': out[72],
+      'subtree_com': out[73],
+      '_impl.subtree_linvel': out[74],
+      '_impl.ten_J': out[75],
+      'ten_length': out[76],
+      '_impl.ten_velocity': out[77],
+      '_impl.ten_wrapadr': out[78],
+      '_impl.ten_wrapnum': out[79],
+      '_impl.tree_island': out[80],
+      '_impl.wrap_obj': out[81],
+      '_impl.wrap_xpos': out[82],
+      'xanchor': out[83],
+      'xaxis': out[84],
+      'ximat': out[85],
+      'xipos': out[86],
+      'xmat': out[87],
+      'xpos': out[88],
+      'xquat': out[89],
+      '_impl.contact__dim': out[90],
+      '_impl.contact__dist': out[91],
+      '_impl.contact__efc_address': out[92],
+      '_impl.contact__flex': out[93],
+      '_impl.contact__frame': out[94],
+      '_impl.contact__friction': out[95],
+      '_impl.contact__geom': out[96],
+      '_impl.contact__geomcollisionid': out[97],
+      '_impl.contact__includemargin': out[98],
+      '_impl.contact__pos': out[99],
+      '_impl.contact__solimp': out[100],
+      '_impl.contact__solref': out[101],
+      '_impl.contact__solreffriction': out[102],
+      '_impl.contact__type': out[103],
+      '_impl.contact__vert': out[104],
+      '_impl.contact__worldid': out[105],
+      '_impl.efc__D': out[106],
+      '_impl.efc__J': out[107],
+      '_impl.efc__J_colind': out[108],
+      '_impl.efc__J_rowadr': out[109],
+      '_impl.efc__J_rownnz': out[110],
+      '_impl.efc__Jqvel': out[111],
+      '_impl.efc__Ma': out[112],
+      '_impl.efc__aref': out[113],
+      '_impl.efc__force': out[114],
+      '_impl.efc__frictionloss': out[115],
+      '_impl.efc__iD': out[116],
+      '_impl.efc__iJ': out[117],
+      '_impl.efc__iJ_colind': out[118],
+      '_impl.efc__iJ_rowadr': out[119],
+      '_impl.efc__iJ_rownnz': out[120],
+      '_impl.efc__iaref': out[121],
+      '_impl.efc__id': out[122],
+      '_impl.efc__iforce': out[123],
+      '_impl.efc__ifrictionloss': out[124],
+      '_impl.efc__iid': out[125],
+      '_impl.efc__island': out[126],
+      '_impl.efc__istate': out[127],
+      '_impl.efc__itype': out[128],
+      '_impl.efc__margin': out[129],
+      '_impl.efc__pos': out[130],
+      '_impl.efc__state': out[131],
+      '_impl.efc__type': out[132],
+      '_impl.efc__vel': out[133],
+  })
+  return d
+
+
+@jax.custom_batching.custom_vmap
+@ffi.marshal_jax_warp_callable
+def forward(m: types.Model, d: types.Data):
+  return _forward_jax_impl(m, d)
+
+
+@forward.def_vmap
+@ffi.marshal_custom_vmap
+def forward_vmap(unused_axis_size, is_batched, m: types.Model, d: types.Data):
+  d = forward(m, d)
+  return d, is_batched[1]
+
+
+@ffi.format_args_for_warp
+def _step_shim(
+    # Model
+    nworld: int,
+    D_colind: wp.array[int],
+    D_diag: wp.array[int],
+    D_rowadr: wp.array[int],
+    D_rownnz: wp.array[int],
+    M_elemid: wp.array2d[int],
+    M_fullm_i: wp.array[int],
+    M_fullm_j: wp.array[int],
+    M_fullm_upper_elemid: wp.array[int],
+    M_fullm_upper_i: wp.array[int],
+    M_fullm_upper_j: wp.array[int],
+    M_mulm_col: wp.array[int],
+    M_mulm_madr: wp.array[int],
+    M_mulm_rowadr: wp.array[int],
+    M_rowadr: wp.array[int],
+    M_rownnz: wp.array[int],
+    M_tiles: tuple[mjwp_types.TileSet, ...],
+    actuator_acc0: wp.array2d[float],
+    actuator_actadr: wp.array[int],
+    actuator_actearly: wp.array[bool],
+    actuator_actlimited: wp.array[bool],
+    actuator_actnum: wp.array[int],
+    actuator_actrange: wp.array2d[wp.vec2],
+    actuator_biasprm: wp.array2d[mjwp_types.vec10f],
+    actuator_biastype: wp.array[int],
+    actuator_cranklength: wp.array2d[float],
+    actuator_ctrllimited: wp.array[bool],
+    actuator_ctrlrange: wp.array2d[wp.vec2],
+    actuator_delay: wp.array[float],
+    actuator_dynprm: wp.array2d[mjwp_types.vec10f],
+    actuator_dyntype: wp.array[int],
+    actuator_forcelimited: wp.array[bool],
+    actuator_forcerange: wp.array2d[wp.vec2],
+    actuator_gainprm: wp.array2d[mjwp_types.vec10f],
+    actuator_gaintype: wp.array[int],
+    actuator_gear: wp.array2d[wp.spatial_vector],
+    actuator_history: wp.array[wp.vec2i],
+    actuator_historyadr: wp.array[int],
+    actuator_lengthrange: wp.array2d[wp.vec2],
+    actuator_trnid: wp.array[wp.vec2i],
+    actuator_trntype: wp.array[int],
+    actuator_trntype_body_adr: wp.array[int],
+    block_dim: mjwp_types.BlockDim,
+    body_branch_start: wp.array[int],
+    body_branches: wp.array[int],
+    body_dofadr: wp.array[int],
+    body_dofnum: wp.array[int],
+    body_fluid_ellipsoid: wp.array[bool],
+    body_geomadr: wp.array[int],
+    body_geomnum: wp.array[int],
+    body_gravcomp: wp.array2d[float],
+    body_inertia: wp.array2d[wp.vec3],
+    body_invweight0: wp.array2d[wp.vec2],
+    body_ipos: wp.array2d[wp.vec3],
+    body_iquat: wp.array2d[wp.quat],
+    body_isdofancestor: wp.array2d[int],
+    body_jntadr: wp.array[int],
+    body_jntnum: wp.array[int],
+    body_mass: wp.array2d[float],
+    body_mocapid: wp.array[int],
+    body_parentid: wp.array[int],
+    body_pos: wp.array2d[wp.vec3],
+    body_quat: wp.array2d[wp.quat],
+    body_rootid: wp.array[int],
+    body_subtreemass: wp.array2d[float],
+    body_tree: tuple[wp.array[int], ...],
+    body_treeid: wp.array[int],
+    body_weldid: wp.array[int],
+    cam_bodyid: wp.array[int],
+    cam_fovy: wp.array2d[float],
+    cam_intrinsic: wp.array2d[wp.vec4],
+    cam_mat0: wp.array2d[wp.mat33],
+    cam_mode: wp.array[int],
+    cam_pos: wp.array2d[wp.vec3],
+    cam_pos0: wp.array2d[wp.vec3],
+    cam_poscom0: wp.array2d[wp.vec3],
+    cam_quat: wp.array2d[wp.quat],
+    cam_resolution: wp.array[wp.vec2i],
+    cam_sensorsize: wp.array[wp.vec2],
+    cam_targetbodyid: wp.array[int],
+    dof_armature: wp.array2d[float],
+    dof_bodyid: wp.array[int],
+    dof_damping: wp.array2d[float],
+    dof_dampingpoly: wp.array2d[wp.vec2],
+    dof_frictionloss: wp.array2d[float],
+    dof_invweight0: wp.array2d[float],
+    dof_jntid: wp.array[int],
+    dof_parentid: wp.array[int],
+    dof_solimp: wp.array2d[mjwp_types.vec5],
+    dof_solref: wp.array2d[wp.vec2],
+    dof_treeid: wp.array[int],
+    dof_tri_col: wp.array[int],
+    dof_tri_row: wp.array[int],
+    eq_connect_adr: wp.array[int],
+    eq_data: wp.array2d[mjwp_types.vec11],
+    eq_flex_adr: wp.array[int],
+    eq_jnt_adr: wp.array[int],
+    eq_obj1id: wp.array[int],
+    eq_obj2id: wp.array[int],
+    eq_objtype: wp.array[int],
+    eq_solimp: wp.array2d[mjwp_types.vec5],
+    eq_solref: wp.array2d[wp.vec2],
+    eq_ten_adr: wp.array[int],
+    eq_type: wp.array[int],
+    eq_wld_adr: wp.array[int],
+    flex_bending: wp.array[float],
+    flex_bendingadr: wp.array[int],
+    flex_centered: wp.array[bool],
+    flex_conaffinity: wp.array[int],
+    flex_condim: wp.array[int],
+    flex_contype: wp.array[int],
+    flex_damping: wp.array[float],
+    flex_dim: wp.array[int],
+    flex_edge: wp.array[wp.vec2i],
+    flex_edgeadr: wp.array[int],
+    flex_edgeflap: wp.array[wp.vec2i],
+    flex_edgenum: wp.array[int],
+    flex_elem: wp.array[int],
+    flex_elemadr: wp.array[int],
+    flex_elemdataadr: wp.array[int],
+    flex_elemedge: wp.array[int],
+    flex_elemedgeadr: wp.array[int],
+    flex_elemnum: wp.array[int],
+    flex_friction: wp.array[wp.vec3],
+    flex_gap: wp.array[float],
+    flex_margin: wp.array[float],
+    flex_priority: wp.array[int],
+    flex_radius: wp.array[float],
+    flex_shell: wp.array[int],
+    flex_shelldataadr: wp.array[int],
+    flex_shellnum: wp.array[int],
+    flex_solimp: wp.array[mjwp_types.vec5],
+    flex_solmix: wp.array[float],
+    flex_solref: wp.array[wp.vec2],
+    flex_stiffness: wp.array[float],
+    flex_stiffnessadr: wp.array[int],
+    flex_vert: wp.array[wp.vec3],
+    flex_vertadr: wp.array[int],
+    flex_vertbodyid: wp.array[int],
+    flex_vertflexid: wp.array[int],
+    flex_vertnum: wp.array[int],
+    flexedge_J_colind: wp.array[int],
+    flexedge_J_rowadr: wp.array[int],
+    flexedge_J_rownnz: wp.array[int],
+    flexedge_invweight0: wp.array[float],
+    flexedge_length0: wp.array[float],
+    geom_aabb: wp.array3d[wp.vec3],
+    geom_bodyid: wp.array[int],
+    geom_conaffinity: wp.array[int],
+    geom_condim: wp.array[int],
+    geom_contype: wp.array[int],
+    geom_dataid: wp.array2d[int],
+    geom_fluid: wp.array2d[float],
+    geom_friction: wp.array2d[wp.vec3],
+    geom_gap: wp.array2d[float],
+    geom_group: wp.array[int],
+    geom_margin: wp.array2d[float],
+    geom_matid: wp.array2d[int],
+    geom_pair_type_count: tuple[int, ...],
+    geom_plugin_index: wp.array[int],
+    geom_pos: wp.array2d[wp.vec3],
+    geom_priority: wp.array[int],
+    geom_quat: wp.array2d[wp.quat],
+    geom_rbound: wp.array2d[float],
+    geom_rgba: wp.array2d[wp.vec4],
+    geom_size: wp.array2d[wp.vec3],
+    geom_solimp: wp.array2d[mjwp_types.vec5],
+    geom_solmix: wp.array2d[float],
+    geom_solref: wp.array2d[wp.vec2],
+    geom_type: wp.array[int],
+    has_fluid: bool,
+    has_sdf_geom: bool,
+    hfield_adr: wp.array[int],
+    hfield_data: wp.array[float],
+    hfield_ncol: wp.array[int],
+    hfield_nrow: wp.array[int],
+    hfield_size: wp.array[wp.vec4],
+    is_sparse: bool,
+    jnt_actfrclimited: wp.array[bool],
+    jnt_actfrcrange: wp.array2d[wp.vec2],
+    jnt_actgravcomp: wp.array[int],
+    jnt_axis: wp.array2d[wp.vec3],
+    jnt_bodyid: wp.array[int],
+    jnt_dofadr: wp.array[int],
+    jnt_limited_ball_adr: wp.array[int],
+    jnt_limited_slide_hinge_adr: wp.array[int],
+    jnt_margin: wp.array2d[float],
+    jnt_pos: wp.array2d[wp.vec3],
+    jnt_qposadr: wp.array[int],
+    jnt_range: wp.array2d[wp.vec2],
+    jnt_solimp: wp.array2d[mjwp_types.vec5],
+    jnt_solref: wp.array2d[wp.vec2],
+    jnt_stiffness: wp.array2d[float],
+    jnt_stiffnesspoly: wp.array2d[wp.vec2],
+    jnt_type: wp.array[int],
+    light_bodyid: wp.array[int],
+    light_dir: wp.array2d[wp.vec3],
+    light_dir0: wp.array2d[wp.vec3],
+    light_mode: wp.array[int],
+    light_pos: wp.array2d[wp.vec3],
+    light_pos0: wp.array2d[wp.vec3],
+    light_poscom0: wp.array2d[wp.vec3],
+    light_targetbodyid: wp.array[int],
+    mapM2D: wp.array[int],
+    mat_rgba: wp.array2d[wp.vec4],
+    max_ten_J_rownnz: int,
+    mesh_face: wp.array[wp.vec3i],
+    mesh_faceadr: wp.array[int],
+    mesh_graph: wp.array[int],
+    mesh_graphadr: wp.array[int],
+    mesh_normal: wp.array[wp.vec3],
+    mesh_normaladr: wp.array[int],
+    mesh_normalnum: wp.array[int],
+    mesh_octadr: wp.array[int],
+    mesh_polyadr: wp.array[int],
+    mesh_polymap: wp.array[int],
+    mesh_polymapadr: wp.array[int],
+    mesh_polymapnum: wp.array[int],
+    mesh_polynormal: wp.array[wp.vec3],
+    mesh_polynum: wp.array[int],
+    mesh_polyvert: wp.array[int],
+    mesh_polyvertadr: wp.array[int],
+    mesh_polyvertnum: wp.array[int],
+    mesh_quat: wp.array[wp.quat],
+    mesh_vert: wp.array[wp.vec3],
+    mesh_vertadr: wp.array[int],
+    mesh_vertnum: wp.array[int],
+    nC: int,
+    nD: int,
+    nJten: int,
+    na: int,
+    nacttrnbody: int,
+    nbody: int,
+    nbranch: int,
+    ncam: int,
+    neq: int,
+    nflex: int,
+    nflexedge: int,
+    nflexelem: int,
+    nflexshelldata: int,
+    nflexvert: int,
+    ngeom: int,
+    ngravcomp: int,
+    nhistory: int,
+    njnt: int,
+    nlight: int,
+    nmaxcondim: int,
+    nmaxmeshdeg: int,
+    nmaxpolygon: int,
+    nmaxpyramid: int,
+    nmeshface: int,
+    nrangefinder: int,
+    nsensorcollision: int,
+    nsensorcontact: int,
+    nsensortaxel: int,
+    nsite: int,
+    ntendon: int,
+    ntree: int,
+    nu: int,
+    nv: int,
+    nv_pad: int,
+    nwrap: int,
+    nxn_geom_pair_filtered: wp.array[wp.vec2i],
+    nxn_pairid: wp.array[wp.vec2i],
+    nxn_pairid_filtered: wp.array[wp.vec2i],
+    oct_aabb: wp.array2d[wp.vec3],
+    oct_child: wp.array[mjwp_types.vec8i],
+    oct_coeff: wp.array[mjwp_types.vec8],
+    pair_dim: wp.array[int],
+    pair_friction: wp.array2d[mjwp_types.vec5],
+    pair_gap: wp.array2d[float],
+    pair_margin: wp.array2d[float],
+    pair_solimp: wp.array2d[mjwp_types.vec5],
+    pair_solref: wp.array2d[wp.vec2],
+    pair_solreffriction: wp.array2d[wp.vec2],
+    plugin: wp.array[int],
+    plugin_attr: wp.array[mjwp_types.vec_pluginattr],
+    qD_fullm_i: wp.array[int],
+    qD_fullm_j: wp.array[int],
+    qLD_all_updates: wp.array[wp.vec3i],
+    qLD_level_offsets: wp.array[int],
+    qLD_updates: tuple[wp.array[wp.vec3i], ...],
+    qpos0: wp.array2d[float],
+    qpos_spring: wp.array2d[float],
+    rangefinder_sensor_adr: wp.array[int],
+    sensor_acc_adr: wp.array[int],
+    sensor_adr: wp.array[int],
+    sensor_adr_to_contact_adr: wp.array[int],
+    sensor_contact_adr: wp.array[int],
+    sensor_cutoff: wp.array[float],
+    sensor_datatype: wp.array[int],
+    sensor_delay: wp.array[float],
+    sensor_dim: wp.array[int],
+    sensor_e_kinetic: bool,
+    sensor_e_potential: bool,
+    sensor_history: wp.array[wp.vec2i],
+    sensor_historyadr: wp.array[int],
+    sensor_interval: wp.array[wp.vec2],
+    sensor_intprm: wp.array2d[int],
+    sensor_limitfrc_adr: wp.array[int],
+    sensor_limitpos_adr: wp.array[int],
+    sensor_limitvel_adr: wp.array[int],
+    sensor_objid: wp.array[int],
+    sensor_objtype: wp.array[int],
+    sensor_pos_adr: wp.array[int],
+    sensor_rangefinder_adr: wp.array[int],
+    sensor_rangefinder_bodyid: wp.array[int],
+    sensor_refid: wp.array[int],
+    sensor_reftype: wp.array[int],
+    sensor_rne_postconstraint: bool,
+    sensor_subtree_vel: bool,
+    sensor_tendonactfrc_adr: wp.array[int],
+    sensor_touch_adr: wp.array[int],
+    sensor_type: wp.array[int],
+    sensor_vel_adr: wp.array[int],
+    site_bodyid: wp.array[int],
+    site_pos: wp.array2d[wp.vec3],
+    site_quat: wp.array2d[wp.quat],
+    site_size: wp.array[wp.vec3],
+    site_type: wp.array[int],
+    taxel_sensorid: wp.array[int],
+    taxel_vertadr: wp.array[int],
+    ten_J_colind: wp.array[int],
+    ten_J_rowadr: wp.array[int],
+    ten_J_rownnz: wp.array[int],
+    tendon_actfrclimited: wp.array[bool],
+    tendon_actfrcrange: wp.array2d[wp.vec2],
+    tendon_adr: wp.array[int],
+    tendon_armature: wp.array2d[float],
+    tendon_damping: wp.array2d[float],
+    tendon_dampingpoly: wp.array2d[wp.vec2],
+    tendon_frictionloss: wp.array2d[float],
+    tendon_geom_adr: wp.array[int],
+    tendon_invweight0: wp.array2d[float],
+    tendon_jnt_adr: wp.array[int],
+    tendon_length0: wp.array2d[float],
+    tendon_lengthspring: wp.array2d[wp.vec2],
+    tendon_limited_adr: wp.array[int],
+    tendon_margin: wp.array2d[float],
+    tendon_num: wp.array[int],
+    tendon_range: wp.array2d[wp.vec2],
+    tendon_site_pair_adr: wp.array[int],
+    tendon_solimp_fri: wp.array2d[mjwp_types.vec5],
+    tendon_solimp_lim: wp.array2d[mjwp_types.vec5],
+    tendon_solref_fri: wp.array2d[wp.vec2],
+    tendon_solref_lim: wp.array2d[wp.vec2],
+    tendon_stiffness: wp.array2d[float],
+    tendon_stiffnesspoly: wp.array2d[wp.vec2],
+    wrap_geom_adr: wp.array[int],
+    wrap_jnt_adr: wp.array[int],
+    wrap_objid: wp.array[int],
+    wrap_prm: wp.array[float],
+    wrap_pulley_scale: wp.array[float],
+    wrap_site_pair_adr: wp.array[int],
+    wrap_type: wp.array[int],
+    opt__broadphase: int,
+    opt__broadphase_filter: int,
+    opt__ccd_iterations: int,
+    opt__ccd_tolerance: wp.array[float],
+    opt__cone: int,
+    opt__contact_sensor_maxmatch: int,
+    opt__density: wp.array[float],
+    opt__disableflags: int,
+    opt__enableflags: int,
+    opt__graph_conditional: bool,
+    opt__gravity: wp.array[wp.vec3],
+    opt__impratio_invsqrt: wp.array[float],
+    opt__integrator: int,
+    opt__iterations: int,
+    opt__ls_iterations: int,
+    opt__ls_parallel: bool,
+    opt__ls_parallel_min_step: float,
+    opt__ls_tolerance: wp.array[float],
+    opt__magnetic: wp.array[wp.vec3],
+    opt__run_collision_detection: bool,
+    opt__sdf_initpoints: int,
+    opt__sdf_iterations: int,
+    opt__solver: int,
+    opt__timestep: wp.array[float],
+    opt__tolerance: wp.array[float],
+    opt__viscosity: wp.array[float],
+    opt__wind: wp.array[wp.vec3],
+    stat__meaninertia: wp.array[float],
+    # Data
+    naccdmax: int,
+    naconmax: int,
+    njmax: int,
+    njmax_nnz: int,
+    M: wp.array3d[float],
+    act: wp.array2d[float],
+    act_dot: wp.array2d[float],
+    actuator_force: wp.array2d[float],
+    actuator_length: wp.array2d[float],
+    actuator_moment: wp.array2d[float],
+    actuator_velocity: wp.array2d[float],
+    cacc: wp.array2d[wp.spatial_vector],
+    cam_xmat: wp.array2d[wp.mat33],
+    cam_xpos: wp.array2d[wp.vec3],
+    cdof: wp.array2d[wp.spatial_vector],
+    cdof_dot: wp.array2d[wp.spatial_vector],
+    cfrc_ext: wp.array2d[wp.spatial_vector],
+    cfrc_int: wp.array2d[wp.spatial_vector],
+    cinert: wp.array2d[mjwp_types.vec10],
+    crb: wp.array2d[mjwp_types.vec10],
+    ctrl: wp.array2d[float],
+    cvel: wp.array2d[wp.spatial_vector],
+    dof_island: wp.array2d[int],
+    dof_islandid: wp.array2d[int],
+    efc_islandid: wp.array2d[int],
+    energy: wp.array[wp.vec2],
+    eq_active: wp.array2d[bool],
+    flexedge_J: wp.array2d[float],
+    flexedge_length: wp.array2d[float],
+    flexedge_velocity: wp.array2d[float],
+    flexvert_xpos: wp.array2d[wp.vec3],
+    geom_xmat: wp.array2d[wp.mat33],
+    geom_xpos: wp.array2d[wp.vec3],
+    history: wp.array2d[float],
+    iqacc: wp.array2d[float],
+    iqacc_smooth: wp.array2d[float],
+    iqfrc_constraint: wp.array2d[float],
+    iqfrc_smooth: wp.array2d[float],
+    island_dofadr: wp.array2d[int],
+    island_efcadr: wp.array2d[int],
+    island_ne: wp.array2d[int],
+    island_nefc: wp.array2d[int],
+    island_nf: wp.array2d[int],
+    island_nv: wp.array2d[int],
+    light_xdir: wp.array2d[wp.vec3],
+    light_xpos: wp.array2d[wp.vec3],
+    map_dof2idof: wp.array2d[int],
+    map_efc2iefc: wp.array2d[int],
+    map_idof2dof: wp.array2d[int],
+    map_iefc2efc: wp.array2d[int],
+    mocap_pos: wp.array2d[wp.vec3],
+    mocap_quat: wp.array2d[wp.quat],
+    moment_colind: wp.array2d[int],
+    moment_rowadr: wp.array2d[int],
+    moment_rownnz: wp.array2d[int],
+    nacon: wp.array[int],
+    ncollision: wp.array[int],
+    ne: wp.array[int],
+    nefc: wp.array[int],
+    nf: wp.array[int],
+    nidof: wp.array[int],
+    nisland: wp.array[int],
+    nl: wp.array[int],
+    qLD: wp.array3d[float],
+    qLDiagInv: wp.array2d[float],
+    qLU: wp.array3d[float],
+    qacc: wp.array2d[float],
+    qacc_smooth: wp.array2d[float],
+    qacc_warmstart: wp.array2d[float],
+    qfrc_actuator: wp.array2d[float],
+    qfrc_applied: wp.array2d[float],
+    qfrc_bias: wp.array2d[float],
+    qfrc_constraint: wp.array2d[float],
+    qfrc_damper: wp.array2d[float],
+    qfrc_fluid: wp.array2d[float],
+    qfrc_gravcomp: wp.array2d[float],
+    qfrc_passive: wp.array2d[float],
+    qfrc_smooth: wp.array2d[float],
+    qfrc_spring: wp.array2d[float],
+    qpos: wp.array2d[float],
+    qvel: wp.array2d[float],
+    sensordata: wp.array2d[float],
+    site_xmat: wp.array2d[wp.mat33],
+    site_xpos: wp.array2d[wp.vec3],
+    solver_niter: wp.array[int],
+    subtree_angmom: wp.array2d[wp.vec3],
+    subtree_com: wp.array2d[wp.vec3],
+    subtree_linvel: wp.array2d[wp.vec3],
+    ten_J: wp.array2d[float],
+    ten_length: wp.array2d[float],
+    ten_velocity: wp.array2d[float],
+    ten_wrapadr: wp.array2d[int],
+    ten_wrapnum: wp.array2d[int],
+    time: wp.array[float],
+    tree_island: wp.array2d[int],
+    wrap_obj: wp.array2d[wp.vec2i],
+    wrap_xpos: wp.array2d[wp.spatial_vector],
+    xanchor: wp.array2d[wp.vec3],
+    xaxis: wp.array2d[wp.vec3],
+    xfrc_applied: wp.array2d[wp.spatial_vector],
+    ximat: wp.array2d[wp.mat33],
+    xipos: wp.array2d[wp.vec3],
+    xmat: wp.array2d[wp.mat33],
+    xpos: wp.array2d[wp.vec3],
+    xquat: wp.array2d[wp.quat],
+    contact__dim: wp.array[int],
+    contact__dist: wp.array[float],
+    contact__efc_address: wp.array2d[int],
+    contact__flex: wp.array[wp.vec2i],
+    contact__frame: wp.array[wp.mat33],
+    contact__friction: wp.array[mjwp_types.vec5],
+    contact__geom: wp.array[wp.vec2i],
+    contact__geomcollisionid: wp.array[int],
+    contact__includemargin: wp.array[float],
+    contact__pos: wp.array[wp.vec3],
+    contact__solimp: wp.array[mjwp_types.vec5],
+    contact__solref: wp.array[wp.vec2],
+    contact__solreffriction: wp.array[wp.vec2],
+    contact__type: wp.array[int],
+    contact__vert: wp.array[wp.vec2i],
+    contact__worldid: wp.array[int],
+    efc__D: wp.array2d[float],
+    efc__J: wp.array3d[float],
+    efc__J_colind: wp.array3d[int],
+    efc__J_rowadr: wp.array2d[int],
+    efc__J_rownnz: wp.array2d[int],
+    efc__Jqvel: wp.array2d[float],
+    efc__Ma: wp.array2d[float],
+    efc__aref: wp.array2d[float],
+    efc__force: wp.array2d[float],
+    efc__frictionloss: wp.array2d[float],
+    efc__iD: wp.array2d[float],
+    efc__iJ: wp.array3d[float],
+    efc__iJ_colind: wp.array3d[int],
+    efc__iJ_rowadr: wp.array2d[int],
+    efc__iJ_rownnz: wp.array2d[int],
+    efc__iaref: wp.array2d[float],
+    efc__id: wp.array2d[int],
+    efc__iforce: wp.array2d[float],
+    efc__ifrictionloss: wp.array2d[float],
+    efc__iid: wp.array2d[int],
+    efc__island: wp.array2d[int],
+    efc__istate: wp.array2d[int],
+    efc__itype: wp.array2d[int],
+    efc__margin: wp.array2d[float],
+    efc__pos: wp.array2d[float],
+    efc__state: wp.array2d[int],
+    efc__type: wp.array2d[int],
+    efc__vel: wp.array2d[float],
+):
+  _m.stat = _s
+  _m.opt = _o
+  _m.callback = _cb
+  _d.efc = _e
+  _d.contact = _c
+  _m.D_colind = D_colind
+  _m.D_diag = D_diag
+  _m.D_rowadr = D_rowadr
+  _m.D_rownnz = D_rownnz
+  _m.M_elemid = M_elemid
+  _m.M_fullm_i = M_fullm_i
+  _m.M_fullm_j = M_fullm_j
+  _m.M_fullm_upper_elemid = M_fullm_upper_elemid
+  _m.M_fullm_upper_i = M_fullm_upper_i
+  _m.M_fullm_upper_j = M_fullm_upper_j
+  _m.M_mulm_col = M_mulm_col
+  _m.M_mulm_madr = M_mulm_madr
+  _m.M_mulm_rowadr = M_mulm_rowadr
+  _m.M_rowadr = M_rowadr
+  _m.M_rownnz = M_rownnz
+  _m.M_tiles = M_tiles
+  _m.actuator_acc0 = actuator_acc0
+  _m.actuator_actadr = actuator_actadr
+  _m.actuator_actearly = actuator_actearly
+  _m.actuator_actlimited = actuator_actlimited
+  _m.actuator_actnum = actuator_actnum
+  _m.actuator_actrange = actuator_actrange
+  _m.actuator_biasprm = actuator_biasprm
+  _m.actuator_biastype = actuator_biastype
+  _m.actuator_cranklength = actuator_cranklength
+  _m.actuator_ctrllimited = actuator_ctrllimited
+  _m.actuator_ctrlrange = actuator_ctrlrange
+  _m.actuator_delay = actuator_delay
+  _m.actuator_dynprm = actuator_dynprm
+  _m.actuator_dyntype = actuator_dyntype
+  _m.actuator_forcelimited = actuator_forcelimited
+  _m.actuator_forcerange = actuator_forcerange
+  _m.actuator_gainprm = actuator_gainprm
+  _m.actuator_gaintype = actuator_gaintype
+  _m.actuator_gear = actuator_gear
+  _m.actuator_history = actuator_history
+  _m.actuator_historyadr = actuator_historyadr
+  _m.actuator_lengthrange = actuator_lengthrange
+  _m.actuator_trnid = actuator_trnid
+  _m.actuator_trntype = actuator_trntype
+  _m.actuator_trntype_body_adr = actuator_trntype_body_adr
+  _m.block_dim = block_dim
+  _m.body_branch_start = body_branch_start
+  _m.body_branches = body_branches
+  _m.body_dofadr = body_dofadr
+  _m.body_dofnum = body_dofnum
+  _m.body_fluid_ellipsoid = body_fluid_ellipsoid
+  _m.body_geomadr = body_geomadr
+  _m.body_geomnum = body_geomnum
+  _m.body_gravcomp = body_gravcomp
+  _m.body_inertia = body_inertia
+  _m.body_invweight0 = body_invweight0
+  _m.body_ipos = body_ipos
+  _m.body_iquat = body_iquat
+  _m.body_isdofancestor = body_isdofancestor
+  _m.body_jntadr = body_jntadr
+  _m.body_jntnum = body_jntnum
+  _m.body_mass = body_mass
+  _m.body_mocapid = body_mocapid
+  _m.body_parentid = body_parentid
+  _m.body_pos = body_pos
+  _m.body_quat = body_quat
+  _m.body_rootid = body_rootid
+  _m.body_subtreemass = body_subtreemass
+  _m.body_tree = body_tree
+  _m.body_treeid = body_treeid
+  _m.body_weldid = body_weldid
+  _m.cam_bodyid = cam_bodyid
+  _m.cam_fovy = cam_fovy
+  _m.cam_intrinsic = cam_intrinsic
+  _m.cam_mat0 = cam_mat0
+  _m.cam_mode = cam_mode
+  _m.cam_pos = cam_pos
+  _m.cam_pos0 = cam_pos0
+  _m.cam_poscom0 = cam_poscom0
+  _m.cam_quat = cam_quat
+  _m.cam_resolution = cam_resolution
+  _m.cam_sensorsize = cam_sensorsize
+  _m.cam_targetbodyid = cam_targetbodyid
+  _m.dof_armature = dof_armature
+  _m.dof_bodyid = dof_bodyid
+  _m.dof_damping = dof_damping
+  _m.dof_dampingpoly = dof_dampingpoly
+  _m.dof_frictionloss = dof_frictionloss
+  _m.dof_invweight0 = dof_invweight0
+  _m.dof_jntid = dof_jntid
+  _m.dof_parentid = dof_parentid
+  _m.dof_solimp = dof_solimp
+  _m.dof_solref = dof_solref
+  _m.dof_treeid = dof_treeid
+  _m.dof_tri_col = dof_tri_col
+  _m.dof_tri_row = dof_tri_row
+  _m.eq_connect_adr = eq_connect_adr
+  _m.eq_data = eq_data
+  _m.eq_flex_adr = eq_flex_adr
+  _m.eq_jnt_adr = eq_jnt_adr
+  _m.eq_obj1id = eq_obj1id
+  _m.eq_obj2id = eq_obj2id
+  _m.eq_objtype = eq_objtype
+  _m.eq_solimp = eq_solimp
+  _m.eq_solref = eq_solref
+  _m.eq_ten_adr = eq_ten_adr
+  _m.eq_type = eq_type
+  _m.eq_wld_adr = eq_wld_adr
+  _m.flex_bending = flex_bending
+  _m.flex_bendingadr = flex_bendingadr
+  _m.flex_centered = flex_centered
+  _m.flex_conaffinity = flex_conaffinity
+  _m.flex_condim = flex_condim
+  _m.flex_contype = flex_contype
+  _m.flex_damping = flex_damping
+  _m.flex_dim = flex_dim
+  _m.flex_edge = flex_edge
+  _m.flex_edgeadr = flex_edgeadr
+  _m.flex_edgeflap = flex_edgeflap
+  _m.flex_edgenum = flex_edgenum
+  _m.flex_elem = flex_elem
+  _m.flex_elemadr = flex_elemadr
+  _m.flex_elemdataadr = flex_elemdataadr
+  _m.flex_elemedge = flex_elemedge
+  _m.flex_elemedgeadr = flex_elemedgeadr
+  _m.flex_elemnum = flex_elemnum
+  _m.flex_friction = flex_friction
+  _m.flex_gap = flex_gap
+  _m.flex_margin = flex_margin
+  _m.flex_priority = flex_priority
+  _m.flex_radius = flex_radius
+  _m.flex_shell = flex_shell
+  _m.flex_shelldataadr = flex_shelldataadr
+  _m.flex_shellnum = flex_shellnum
+  _m.flex_solimp = flex_solimp
+  _m.flex_solmix = flex_solmix
+  _m.flex_solref = flex_solref
+  _m.flex_stiffness = flex_stiffness
+  _m.flex_stiffnessadr = flex_stiffnessadr
+  _m.flex_vert = flex_vert
+  _m.flex_vertadr = flex_vertadr
+  _m.flex_vertbodyid = flex_vertbodyid
+  _m.flex_vertflexid = flex_vertflexid
+  _m.flex_vertnum = flex_vertnum
+  _m.flexedge_J_colind = flexedge_J_colind
+  _m.flexedge_J_rowadr = flexedge_J_rowadr
+  _m.flexedge_J_rownnz = flexedge_J_rownnz
+  _m.flexedge_invweight0 = flexedge_invweight0
+  _m.flexedge_length0 = flexedge_length0
+  _m.geom_aabb = geom_aabb
+  _m.geom_bodyid = geom_bodyid
+  _m.geom_conaffinity = geom_conaffinity
+  _m.geom_condim = geom_condim
+  _m.geom_contype = geom_contype
+  _m.geom_dataid = geom_dataid
+  _m.geom_fluid = geom_fluid
+  _m.geom_friction = geom_friction
+  _m.geom_gap = geom_gap
+  _m.geom_group = geom_group
+  _m.geom_margin = geom_margin
+  _m.geom_matid = geom_matid
+  _m.geom_pair_type_count = geom_pair_type_count
+  _m.geom_plugin_index = geom_plugin_index
+  _m.geom_pos = geom_pos
+  _m.geom_priority = geom_priority
+  _m.geom_quat = geom_quat
+  _m.geom_rbound = geom_rbound
+  _m.geom_rgba = geom_rgba
+  _m.geom_size = geom_size
+  _m.geom_solimp = geom_solimp
+  _m.geom_solmix = geom_solmix
+  _m.geom_solref = geom_solref
+  _m.geom_type = geom_type
+  _m.has_fluid = has_fluid
+  _m.has_sdf_geom = has_sdf_geom
+  _m.hfield_adr = hfield_adr
+  _m.hfield_data = hfield_data
+  _m.hfield_ncol = hfield_ncol
+  _m.hfield_nrow = hfield_nrow
+  _m.hfield_size = hfield_size
+  _m.is_sparse = is_sparse
+  _m.jnt_actfrclimited = jnt_actfrclimited
+  _m.jnt_actfrcrange = jnt_actfrcrange
+  _m.jnt_actgravcomp = jnt_actgravcomp
+  _m.jnt_axis = jnt_axis
+  _m.jnt_bodyid = jnt_bodyid
+  _m.jnt_dofadr = jnt_dofadr
+  _m.jnt_limited_ball_adr = jnt_limited_ball_adr
+  _m.jnt_limited_slide_hinge_adr = jnt_limited_slide_hinge_adr
+  _m.jnt_margin = jnt_margin
+  _m.jnt_pos = jnt_pos
+  _m.jnt_qposadr = jnt_qposadr
+  _m.jnt_range = jnt_range
+  _m.jnt_solimp = jnt_solimp
+  _m.jnt_solref = jnt_solref
+  _m.jnt_stiffness = jnt_stiffness
+  _m.jnt_stiffnesspoly = jnt_stiffnesspoly
+  _m.jnt_type = jnt_type
+  _m.light_bodyid = light_bodyid
+  _m.light_dir = light_dir
+  _m.light_dir0 = light_dir0
+  _m.light_mode = light_mode
+  _m.light_pos = light_pos
+  _m.light_pos0 = light_pos0
+  _m.light_poscom0 = light_poscom0
+  _m.light_targetbodyid = light_targetbodyid
+  _m.mapM2D = mapM2D
+  _m.mat_rgba = mat_rgba
+  _m.max_ten_J_rownnz = max_ten_J_rownnz
+  _m.mesh_face = mesh_face
+  _m.mesh_faceadr = mesh_faceadr
+  _m.mesh_graph = mesh_graph
+  _m.mesh_graphadr = mesh_graphadr
+  _m.mesh_normal = mesh_normal
+  _m.mesh_normaladr = mesh_normaladr
+  _m.mesh_normalnum = mesh_normalnum
+  _m.mesh_octadr = mesh_octadr
+  _m.mesh_polyadr = mesh_polyadr
+  _m.mesh_polymap = mesh_polymap
+  _m.mesh_polymapadr = mesh_polymapadr
+  _m.mesh_polymapnum = mesh_polymapnum
+  _m.mesh_polynormal = mesh_polynormal
+  _m.mesh_polynum = mesh_polynum
+  _m.mesh_polyvert = mesh_polyvert
+  _m.mesh_polyvertadr = mesh_polyvertadr
+  _m.mesh_polyvertnum = mesh_polyvertnum
+  _m.mesh_quat = mesh_quat
+  _m.mesh_vert = mesh_vert
+  _m.mesh_vertadr = mesh_vertadr
+  _m.mesh_vertnum = mesh_vertnum
+  _m.nC = nC
+  _m.nD = nD
+  _m.nJten = nJten
+  _m.na = na
+  _m.nacttrnbody = nacttrnbody
+  _m.nbody = nbody
+  _m.nbranch = nbranch
+  _m.ncam = ncam
+  _m.neq = neq
+  _m.nflex = nflex
+  _m.nflexedge = nflexedge
+  _m.nflexelem = nflexelem
+  _m.nflexshelldata = nflexshelldata
+  _m.nflexvert = nflexvert
+  _m.ngeom = ngeom
+  _m.ngravcomp = ngravcomp
+  _m.nhistory = nhistory
+  _m.njnt = njnt
+  _m.nlight = nlight
+  _m.nmaxcondim = nmaxcondim
+  _m.nmaxmeshdeg = nmaxmeshdeg
+  _m.nmaxpolygon = nmaxpolygon
+  _m.nmaxpyramid = nmaxpyramid
+  _m.nmeshface = nmeshface
+  _m.nrangefinder = nrangefinder
+  _m.nsensorcollision = nsensorcollision
+  _m.nsensorcontact = nsensorcontact
+  _m.nsensortaxel = nsensortaxel
+  _m.nsite = nsite
+  _m.ntendon = ntendon
+  _m.ntree = ntree
+  _m.nu = nu
+  _m.nv = nv
+  _m.nv_pad = nv_pad
+  _m.nwrap = nwrap
+  _m.nxn_geom_pair_filtered = nxn_geom_pair_filtered
+  _m.nxn_pairid = nxn_pairid
+  _m.nxn_pairid_filtered = nxn_pairid_filtered
+  _m.oct_aabb = oct_aabb
+  _m.oct_child = oct_child
+  _m.oct_coeff = oct_coeff
+  _m.opt.broadphase = opt__broadphase
+  _m.opt.broadphase_filter = opt__broadphase_filter
+  _m.opt.ccd_iterations = opt__ccd_iterations
+  _m.opt.ccd_tolerance = opt__ccd_tolerance
+  _m.opt.cone = opt__cone
+  _m.opt.contact_sensor_maxmatch = opt__contact_sensor_maxmatch
+  _m.opt.density = opt__density
+  _m.opt.disableflags = opt__disableflags
+  _m.opt.enableflags = opt__enableflags
+  _m.opt.graph_conditional = opt__graph_conditional
+  _m.opt.gravity = opt__gravity
+  _m.opt.impratio_invsqrt = opt__impratio_invsqrt
+  _m.opt.integrator = opt__integrator
+  _m.opt.iterations = opt__iterations
+  _m.opt.ls_iterations = opt__ls_iterations
+  _m.opt.ls_parallel = opt__ls_parallel
+  _m.opt.ls_parallel_min_step = opt__ls_parallel_min_step
+  _m.opt.ls_tolerance = opt__ls_tolerance
+  _m.opt.magnetic = opt__magnetic
+  _m.opt.run_collision_detection = opt__run_collision_detection
+  _m.opt.sdf_initpoints = opt__sdf_initpoints
+  _m.opt.sdf_iterations = opt__sdf_iterations
+  _m.opt.solver = opt__solver
+  _m.opt.timestep = opt__timestep
+  _m.opt.tolerance = opt__tolerance
+  _m.opt.viscosity = opt__viscosity
+  _m.opt.wind = opt__wind
+  _m.pair_dim = pair_dim
+  _m.pair_friction = pair_friction
+  _m.pair_gap = pair_gap
+  _m.pair_margin = pair_margin
+  _m.pair_solimp = pair_solimp
+  _m.pair_solref = pair_solref
+  _m.pair_solreffriction = pair_solreffriction
+  _m.plugin = plugin
+  _m.plugin_attr = plugin_attr
+  _m.qD_fullm_i = qD_fullm_i
+  _m.qD_fullm_j = qD_fullm_j
+  _m.qLD_all_updates = qLD_all_updates
+  _m.qLD_level_offsets = qLD_level_offsets
+  _m.qLD_updates = qLD_updates
+  _m.qpos0 = qpos0
+  _m.qpos_spring = qpos_spring
+  _m.rangefinder_sensor_adr = rangefinder_sensor_adr
+  _m.sensor_acc_adr = sensor_acc_adr
+  _m.sensor_adr = sensor_adr
+  _m.sensor_adr_to_contact_adr = sensor_adr_to_contact_adr
+  _m.sensor_contact_adr = sensor_contact_adr
+  _m.sensor_cutoff = sensor_cutoff
+  _m.sensor_datatype = sensor_datatype
+  _m.sensor_delay = sensor_delay
+  _m.sensor_dim = sensor_dim
+  _m.sensor_e_kinetic = sensor_e_kinetic
+  _m.sensor_e_potential = sensor_e_potential
+  _m.sensor_history = sensor_history
+  _m.sensor_historyadr = sensor_historyadr
+  _m.sensor_interval = sensor_interval
+  _m.sensor_intprm = sensor_intprm
+  _m.sensor_limitfrc_adr = sensor_limitfrc_adr
+  _m.sensor_limitpos_adr = sensor_limitpos_adr
+  _m.sensor_limitvel_adr = sensor_limitvel_adr
+  _m.sensor_objid = sensor_objid
+  _m.sensor_objtype = sensor_objtype
+  _m.sensor_pos_adr = sensor_pos_adr
+  _m.sensor_rangefinder_adr = sensor_rangefinder_adr
+  _m.sensor_rangefinder_bodyid = sensor_rangefinder_bodyid
+  _m.sensor_refid = sensor_refid
+  _m.sensor_reftype = sensor_reftype
+  _m.sensor_rne_postconstraint = sensor_rne_postconstraint
+  _m.sensor_subtree_vel = sensor_subtree_vel
+  _m.sensor_tendonactfrc_adr = sensor_tendonactfrc_adr
+  _m.sensor_touch_adr = sensor_touch_adr
+  _m.sensor_type = sensor_type
+  _m.sensor_vel_adr = sensor_vel_adr
+  _m.site_bodyid = site_bodyid
+  _m.site_pos = site_pos
+  _m.site_quat = site_quat
+  _m.site_size = site_size
+  _m.site_type = site_type
+  _m.stat.meaninertia = stat__meaninertia
+  _m.taxel_sensorid = taxel_sensorid
+  _m.taxel_vertadr = taxel_vertadr
+  _m.ten_J_colind = ten_J_colind
+  _m.ten_J_rowadr = ten_J_rowadr
+  _m.ten_J_rownnz = ten_J_rownnz
+  _m.tendon_actfrclimited = tendon_actfrclimited
+  _m.tendon_actfrcrange = tendon_actfrcrange
+  _m.tendon_adr = tendon_adr
+  _m.tendon_armature = tendon_armature
+  _m.tendon_damping = tendon_damping
+  _m.tendon_dampingpoly = tendon_dampingpoly
+  _m.tendon_frictionloss = tendon_frictionloss
+  _m.tendon_geom_adr = tendon_geom_adr
+  _m.tendon_invweight0 = tendon_invweight0
+  _m.tendon_jnt_adr = tendon_jnt_adr
+  _m.tendon_length0 = tendon_length0
+  _m.tendon_lengthspring = tendon_lengthspring
+  _m.tendon_limited_adr = tendon_limited_adr
+  _m.tendon_margin = tendon_margin
+  _m.tendon_num = tendon_num
+  _m.tendon_range = tendon_range
+  _m.tendon_site_pair_adr = tendon_site_pair_adr
+  _m.tendon_solimp_fri = tendon_solimp_fri
+  _m.tendon_solimp_lim = tendon_solimp_lim
+  _m.tendon_solref_fri = tendon_solref_fri
+  _m.tendon_solref_lim = tendon_solref_lim
+  _m.tendon_stiffness = tendon_stiffness
+  _m.tendon_stiffnesspoly = tendon_stiffnesspoly
+  _m.wrap_geom_adr = wrap_geom_adr
+  _m.wrap_jnt_adr = wrap_jnt_adr
+  _m.wrap_objid = wrap_objid
+  _m.wrap_prm = wrap_prm
+  _m.wrap_pulley_scale = wrap_pulley_scale
+  _m.wrap_site_pair_adr = wrap_site_pair_adr
+  _m.wrap_type = wrap_type
+  _d.M = M
+  _d.act = act
+  _d.act_dot = act_dot
+  _d.actuator_force = actuator_force
+  _d.actuator_length = actuator_length
+  _d.actuator_moment = actuator_moment
+  _d.actuator_velocity = actuator_velocity
+  _d.cacc = cacc
+  _d.cam_xmat = cam_xmat
+  _d.cam_xpos = cam_xpos
+  _d.cdof = cdof
+  _d.cdof_dot = cdof_dot
+  _d.cfrc_ext = cfrc_ext
+  _d.cfrc_int = cfrc_int
+  _d.cinert = cinert
+  _d.contact.dim = contact__dim
+  _d.contact.dist = contact__dist
+  _d.contact.efc_address = contact__efc_address
+  _d.contact.flex = contact__flex
+  _d.contact.frame = contact__frame
+  _d.contact.friction = contact__friction
+  _d.contact.geom = contact__geom
+  _d.contact.geomcollisionid = contact__geomcollisionid
+  _d.contact.includemargin = contact__includemargin
+  _d.contact.pos = contact__pos
+  _d.contact.solimp = contact__solimp
+  _d.contact.solref = contact__solref
+  _d.contact.solreffriction = contact__solreffriction
+  _d.contact.type = contact__type
+  _d.contact.vert = contact__vert
+  _d.contact.worldid = contact__worldid
+  _d.crb = crb
+  _d.ctrl = ctrl
+  _d.cvel = cvel
+  _d.dof_island = dof_island
+  _d.dof_islandid = dof_islandid
+  _d.efc.D = efc__D
+  _d.efc.J = efc__J
+  _d.efc.J_colind = efc__J_colind
+  _d.efc.J_rowadr = efc__J_rowadr
+  _d.efc.J_rownnz = efc__J_rownnz
+  _d.efc.Jqvel = efc__Jqvel
+  _d.efc.Ma = efc__Ma
+  _d.efc.aref = efc__aref
+  _d.efc.force = efc__force
+  _d.efc.frictionloss = efc__frictionloss
+  _d.efc.iD = efc__iD
+  _d.efc.iJ = efc__iJ
+  _d.efc.iJ_colind = efc__iJ_colind
+  _d.efc.iJ_rowadr = efc__iJ_rowadr
+  _d.efc.iJ_rownnz = efc__iJ_rownnz
+  _d.efc.iaref = efc__iaref
+  _d.efc.id = efc__id
+  _d.efc.iforce = efc__iforce
+  _d.efc.ifrictionloss = efc__ifrictionloss
+  _d.efc.iid = efc__iid
+  _d.efc.island = efc__island
+  _d.efc.istate = efc__istate
+  _d.efc.itype = efc__itype
+  _d.efc.margin = efc__margin
+  _d.efc.pos = efc__pos
+  _d.efc.state = efc__state
+  _d.efc.type = efc__type
+  _d.efc.vel = efc__vel
+  _d.efc_islandid = efc_islandid
+  _d.energy = energy
+  _d.eq_active = eq_active
+  _d.flexedge_J = flexedge_J
+  _d.flexedge_length = flexedge_length
+  _d.flexedge_velocity = flexedge_velocity
+  _d.flexvert_xpos = flexvert_xpos
+  _d.geom_xmat = geom_xmat
+  _d.geom_xpos = geom_xpos
+  _d.history = history
+  _d.iqacc = iqacc
+  _d.iqacc_smooth = iqacc_smooth
+  _d.iqfrc_constraint = iqfrc_constraint
+  _d.iqfrc_smooth = iqfrc_smooth
+  _d.island_dofadr = island_dofadr
+  _d.island_efcadr = island_efcadr
+  _d.island_ne = island_ne
+  _d.island_nefc = island_nefc
+  _d.island_nf = island_nf
+  _d.island_nv = island_nv
+  _d.light_xdir = light_xdir
+  _d.light_xpos = light_xpos
+  _d.map_dof2idof = map_dof2idof
+  _d.map_efc2iefc = map_efc2iefc
+  _d.map_idof2dof = map_idof2dof
+  _d.map_iefc2efc = map_iefc2efc
+  _d.mocap_pos = mocap_pos
+  _d.mocap_quat = mocap_quat
+  _d.moment_colind = moment_colind
+  _d.moment_rowadr = moment_rowadr
+  _d.moment_rownnz = moment_rownnz
+  _d.naccdmax = naccdmax
+  _d.nacon = nacon
+  _d.naconmax = naconmax
+  _d.ncollision = ncollision
+  _d.ne = ne
+  _d.nefc = nefc
+  _d.nf = nf
+  _d.nidof = nidof
+  _d.nisland = nisland
+  _d.njmax = njmax
+  _d.njmax_nnz = njmax_nnz
+  _d.nl = nl
+  _d.qLD = qLD
+  _d.qLDiagInv = qLDiagInv
+  _d.qLU = qLU
+  _d.qacc = qacc
+  _d.qacc_smooth = qacc_smooth
+  _d.qacc_warmstart = qacc_warmstart
+  _d.qfrc_actuator = qfrc_actuator
+  _d.qfrc_applied = qfrc_applied
+  _d.qfrc_bias = qfrc_bias
+  _d.qfrc_constraint = qfrc_constraint
+  _d.qfrc_damper = qfrc_damper
+  _d.qfrc_fluid = qfrc_fluid
+  _d.qfrc_gravcomp = qfrc_gravcomp
+  _d.qfrc_passive = qfrc_passive
+  _d.qfrc_smooth = qfrc_smooth
+  _d.qfrc_spring = qfrc_spring
+  _d.qpos = qpos
+  _d.qvel = qvel
+  _d.sensordata = sensordata
+  _d.site_xmat = site_xmat
+  _d.site_xpos = site_xpos
+  _d.solver_niter = solver_niter
+  _d.subtree_angmom = subtree_angmom
+  _d.subtree_com = subtree_com
+  _d.subtree_linvel = subtree_linvel
+  _d.ten_J = ten_J
+  _d.ten_length = ten_length
+  _d.ten_velocity = ten_velocity
+  _d.ten_wrapadr = ten_wrapadr
+  _d.ten_wrapnum = ten_wrapnum
+  _d.time = time
+  _d.tree_island = tree_island
+  _d.wrap_obj = wrap_obj
+  _d.wrap_xpos = wrap_xpos
+  _d.xanchor = xanchor
+  _d.xaxis = xaxis
+  _d.xfrc_applied = xfrc_applied
+  _d.ximat = ximat
+  _d.xipos = xipos
+  _d.xmat = xmat
+  _d.xpos = xpos
+  _d.xquat = xquat
+  _d.nworld = nworld
+  mjwarp.step(_m, _d)
+
+
+def _step_jax_impl(m: types.Model, d: types.Data):
+  output_dims = {
+      'M': d._impl.M.shape,
+      'act': d.act.shape,
+      'act_dot': d.act_dot.shape,
+      'actuator_force': d.actuator_force.shape,
+      'actuator_length': d.actuator_length.shape,
+      'actuator_moment': d._impl.actuator_moment.shape,
+      'actuator_velocity': d._impl.actuator_velocity.shape,
+      'cacc': d._impl.cacc.shape,
+      'cam_xmat': d.cam_xmat.shape,
+      'cam_xpos': d.cam_xpos.shape,
+      'cdof': d.cdof.shape,
+      'cdof_dot': d.cdof_dot.shape,
+      'cfrc_ext': d._impl.cfrc_ext.shape,
+      'cfrc_int': d._impl.cfrc_int.shape,
+      'cinert': d._impl.cinert.shape,
+      'crb': d._impl.crb.shape,
+      'cvel': d.cvel.shape,
+      'dof_island': d._impl.dof_island.shape,
+      'dof_islandid': d._impl.dof_islandid.shape,
+      'efc_islandid': d._impl.efc_islandid.shape,
+      'energy': d._impl.energy.shape,
+      'flexedge_J': d._impl.flexedge_J.shape,
+      'flexedge_length': d._impl.flexedge_length.shape,
+      'flexedge_velocity': d._impl.flexedge_velocity.shape,
+      'flexvert_xpos': d._impl.flexvert_xpos.shape,
+      'geom_xmat': d.geom_xmat.shape,
+      'geom_xpos': d.geom_xpos.shape,
+      'history': d.history.shape,
+      'iqacc': d._impl.iqacc.shape,
+      'iqacc_smooth': d._impl.iqacc_smooth.shape,
+      'iqfrc_constraint': d._impl.iqfrc_constraint.shape,
+      'iqfrc_smooth': d._impl.iqfrc_smooth.shape,
+      'island_dofadr': d._impl.island_dofadr.shape,
+      'island_efcadr': d._impl.island_efcadr.shape,
+      'island_ne': d._impl.island_ne.shape,
+      'island_nefc': d._impl.island_nefc.shape,
+      'island_nf': d._impl.island_nf.shape,
+      'island_nv': d._impl.island_nv.shape,
+      'light_xdir': d._impl.light_xdir.shape,
+      'light_xpos': d._impl.light_xpos.shape,
+      'map_dof2idof': d._impl.map_dof2idof.shape,
+      'map_efc2iefc': d._impl.map_efc2iefc.shape,
+      'map_idof2dof': d._impl.map_idof2dof.shape,
+      'map_iefc2efc': d._impl.map_iefc2efc.shape,
+      'moment_colind': d._impl.moment_colind.shape,
+      'moment_rowadr': d._impl.moment_rowadr.shape,
+      'moment_rownnz': d._impl.moment_rownnz.shape,
+      'nacon': d._impl.nacon.shape,
+      'ncollision': d._impl.ncollision.shape,
+      'ne': d._impl.ne.shape,
+      'nefc': d._impl.nefc.shape,
+      'nf': d._impl.nf.shape,
+      'nidof': d._impl.nidof.shape,
+      'nisland': d._impl.nisland.shape,
+      'nl': d._impl.nl.shape,
+      'qLD': d._impl.qLD.shape,
+      'qLDiagInv': d._impl.qLDiagInv.shape,
+      'qLU': d._impl.qLU.shape,
+      'qacc': d.qacc.shape,
+      'qacc_smooth': d.qacc_smooth.shape,
+      'qacc_warmstart': d.qacc_warmstart.shape,
+      'qfrc_actuator': d.qfrc_actuator.shape,
+      'qfrc_bias': d.qfrc_bias.shape,
+      'qfrc_constraint': d.qfrc_constraint.shape,
+      'qfrc_damper': d._impl.qfrc_damper.shape,
+      'qfrc_fluid': d.qfrc_fluid.shape,
+      'qfrc_gravcomp': d.qfrc_gravcomp.shape,
+      'qfrc_passive': d.qfrc_passive.shape,
+      'qfrc_smooth': d.qfrc_smooth.shape,
+      'qfrc_spring': d._impl.qfrc_spring.shape,
+      'qpos': d.qpos.shape,
+      'qvel': d.qvel.shape,
+      'sensordata': d.sensordata.shape,
+      'site_xmat': d.site_xmat.shape,
+      'site_xpos': d.site_xpos.shape,
+      'solver_niter': d._impl.solver_niter.shape,
+      'subtree_angmom': d._impl.subtree_angmom.shape,
+      'subtree_com': d.subtree_com.shape,
+      'subtree_linvel': d._impl.subtree_linvel.shape,
+      'ten_J': d._impl.ten_J.shape,
+      'ten_length': d.ten_length.shape,
+      'ten_velocity': d._impl.ten_velocity.shape,
+      'ten_wrapadr': d._impl.ten_wrapadr.shape,
+      'ten_wrapnum': d._impl.ten_wrapnum.shape,
+      'time': d.time.shape,
+      'tree_island': d._impl.tree_island.shape,
+      'wrap_obj': d._impl.wrap_obj.shape,
+      'wrap_xpos': d._impl.wrap_xpos.shape,
+      'xanchor': d.xanchor.shape,
+      'xaxis': d.xaxis.shape,
+      'ximat': d.ximat.shape,
+      'xipos': d.xipos.shape,
+      'xmat': d.xmat.shape,
+      'xpos': d.xpos.shape,
+      'xquat': d.xquat.shape,
+      'contact__dim': d._impl.contact__dim.shape,
+      'contact__dist': d._impl.contact__dist.shape,
+      'contact__efc_address': d._impl.contact__efc_address.shape,
+      'contact__flex': d._impl.contact__flex.shape,
+      'contact__frame': d._impl.contact__frame.shape,
+      'contact__friction': d._impl.contact__friction.shape,
+      'contact__geom': d._impl.contact__geom.shape,
+      'contact__geomcollisionid': d._impl.contact__geomcollisionid.shape,
+      'contact__includemargin': d._impl.contact__includemargin.shape,
+      'contact__pos': d._impl.contact__pos.shape,
+      'contact__solimp': d._impl.contact__solimp.shape,
+      'contact__solref': d._impl.contact__solref.shape,
+      'contact__solreffriction': d._impl.contact__solreffriction.shape,
+      'contact__type': d._impl.contact__type.shape,
+      'contact__vert': d._impl.contact__vert.shape,
+      'contact__worldid': d._impl.contact__worldid.shape,
+      'efc__D': d._impl.efc__D.shape,
+      'efc__J': d._impl.efc__J.shape,
+      'efc__J_colind': d._impl.efc__J_colind.shape,
+      'efc__J_rowadr': d._impl.efc__J_rowadr.shape,
+      'efc__J_rownnz': d._impl.efc__J_rownnz.shape,
+      'efc__Jqvel': d._impl.efc__Jqvel.shape,
+      'efc__Ma': d._impl.efc__Ma.shape,
+      'efc__aref': d._impl.efc__aref.shape,
+      'efc__force': d._impl.efc__force.shape,
+      'efc__frictionloss': d._impl.efc__frictionloss.shape,
+      'efc__iD': d._impl.efc__iD.shape,
+      'efc__iJ': d._impl.efc__iJ.shape,
+      'efc__iJ_colind': d._impl.efc__iJ_colind.shape,
+      'efc__iJ_rowadr': d._impl.efc__iJ_rowadr.shape,
+      'efc__iJ_rownnz': d._impl.efc__iJ_rownnz.shape,
+      'efc__iaref': d._impl.efc__iaref.shape,
+      'efc__id': d._impl.efc__id.shape,
+      'efc__iforce': d._impl.efc__iforce.shape,
+      'efc__ifrictionloss': d._impl.efc__ifrictionloss.shape,
+      'efc__iid': d._impl.efc__iid.shape,
+      'efc__island': d._impl.efc__island.shape,
+      'efc__istate': d._impl.efc__istate.shape,
+      'efc__itype': d._impl.efc__itype.shape,
+      'efc__margin': d._impl.efc__margin.shape,
+      'efc__pos': d._impl.efc__pos.shape,
+      'efc__state': d._impl.efc__state.shape,
+      'efc__type': d._impl.efc__type.shape,
+      'efc__vel': d._impl.efc__vel.shape,
+  }
+  jf = ffi.jax_callable_variadic_tuple(
+      _step_shim,
+      num_outputs=139,
+      output_dims=output_dims,
+      vmap_method=None,
+      in_out_argnames=set([
+          'M',
+          'act',
+          'act_dot',
+          'actuator_force',
+          'actuator_length',
+          'actuator_moment',
+          'actuator_velocity',
+          'cacc',
+          'cam_xmat',
+          'cam_xpos',
+          'cdof',
+          'cdof_dot',
+          'cfrc_ext',
+          'cfrc_int',
+          'cinert',
+          'crb',
+          'cvel',
+          'dof_island',
+          'dof_islandid',
+          'efc_islandid',
+          'energy',
+          'flexedge_J',
+          'flexedge_length',
+          'flexedge_velocity',
+          'flexvert_xpos',
+          'geom_xmat',
+          'geom_xpos',
+          'history',
+          'iqacc',
+          'iqacc_smooth',
+          'iqfrc_constraint',
+          'iqfrc_smooth',
+          'island_dofadr',
+          'island_efcadr',
+          'island_ne',
+          'island_nefc',
+          'island_nf',
+          'island_nv',
+          'light_xdir',
+          'light_xpos',
+          'map_dof2idof',
+          'map_efc2iefc',
+          'map_idof2dof',
+          'map_iefc2efc',
+          'moment_colind',
+          'moment_rowadr',
+          'moment_rownnz',
+          'nacon',
+          'ncollision',
+          'ne',
+          'nefc',
+          'nf',
+          'nidof',
+          'nisland',
+          'nl',
+          'qLD',
+          'qLDiagInv',
+          'qLU',
+          'qacc',
+          'qacc_smooth',
+          'qacc_warmstart',
+          'qfrc_actuator',
+          'qfrc_bias',
+          'qfrc_constraint',
+          'qfrc_damper',
+          'qfrc_fluid',
+          'qfrc_gravcomp',
+          'qfrc_passive',
+          'qfrc_smooth',
+          'qfrc_spring',
+          'qpos',
+          'qvel',
+          'sensordata',
+          'site_xmat',
+          'site_xpos',
+          'solver_niter',
+          'subtree_angmom',
+          'subtree_com',
+          'subtree_linvel',
+          'ten_J',
+          'ten_length',
+          'ten_velocity',
+          'ten_wrapadr',
+          'ten_wrapnum',
+          'time',
+          'tree_island',
+          'wrap_obj',
+          'wrap_xpos',
+          'xanchor',
+          'xaxis',
+          'ximat',
+          'xipos',
+          'xmat',
+          'xpos',
+          'xquat',
+          'contact__dim',
+          'contact__dist',
+          'contact__efc_address',
+          'contact__flex',
+          'contact__frame',
+          'contact__friction',
+          'contact__geom',
+          'contact__geomcollisionid',
+          'contact__includemargin',
+          'contact__pos',
+          'contact__solimp',
+          'contact__solref',
+          'contact__solreffriction',
+          'contact__type',
+          'contact__vert',
+          'contact__worldid',
+          'efc__D',
+          'efc__J',
+          'efc__J_colind',
+          'efc__J_rowadr',
+          'efc__J_rownnz',
+          'efc__Jqvel',
+          'efc__Ma',
+          'efc__aref',
+          'efc__force',
+          'efc__frictionloss',
+          'efc__iD',
+          'efc__iJ',
+          'efc__iJ_colind',
+          'efc__iJ_rowadr',
+          'efc__iJ_rownnz',
+          'efc__iaref',
+          'efc__id',
+          'efc__iforce',
+          'efc__ifrictionloss',
+          'efc__iid',
+          'efc__island',
+          'efc__istate',
+          'efc__itype',
+          'efc__margin',
+          'efc__pos',
+          'efc__state',
+          'efc__type',
+          'efc__vel',
+      ]),
+      stage_in_argnames=set([
+          'act',
+          'act_dot',
+          'actuator_acc0',
+          'actuator_actrange',
+          'actuator_biasprm',
+          'actuator_cranklength',
+          'actuator_ctrlrange',
+          'actuator_dynprm',
+          'actuator_force',
+          'actuator_forcerange',
+          'actuator_gainprm',
+          'actuator_gear',
+          'actuator_length',
+          'actuator_lengthrange',
+          'body_gravcomp',
+          'body_inertia',
+          'body_invweight0',
+          'body_ipos',
+          'body_iquat',
+          'body_mass',
+          'body_pos',
+          'body_quat',
+          'body_subtreemass',
+          'cam_fovy',
+          'cam_intrinsic',
+          'cam_mat0',
+          'cam_pos',
+          'cam_pos0',
+          'cam_poscom0',
+          'cam_quat',
+          'cam_xmat',
+          'cam_xpos',
+          'cdof',
+          'cdof_dot',
+          'ctrl',
+          'cvel',
+          'dof_armature',
+          'dof_damping',
+          'dof_dampingpoly',
+          'dof_frictionloss',
+          'dof_invweight0',
+          'dof_solimp',
+          'dof_solref',
+          'eq_active',
+          'eq_data',
+          'eq_solimp',
+          'eq_solref',
+          'geom_aabb',
+          'geom_friction',
+          'geom_gap',
+          'geom_margin',
+          'geom_matid',
+          'geom_pos',
+          'geom_quat',
+          'geom_rbound',
+          'geom_rgba',
+          'geom_size',
+          'geom_solimp',
+          'geom_solmix',
+          'geom_solref',
+          'geom_xmat',
+          'geom_xpos',
+          'hfield_data',
+          'history',
+          'jnt_actfrcrange',
+          'jnt_axis',
+          'jnt_margin',
+          'jnt_pos',
+          'jnt_range',
+          'jnt_solimp',
+          'jnt_solref',
+          'jnt_stiffness',
+          'jnt_stiffnesspoly',
+          'light_dir',
+          'light_dir0',
+          'light_pos',
+          'light_pos0',
+          'light_poscom0',
+          'mat_rgba',
+          'mocap_pos',
+          'mocap_quat',
+          'opt__density',
+          'opt__gravity',
+          'opt__ls_tolerance',
+          'opt__magnetic',
+          'opt__timestep',
+          'opt__tolerance',
+          'opt__viscosity',
+          'opt__wind',
+          'pair_friction',
+          'pair_gap',
+          'pair_margin',
+          'pair_solimp',
+          'pair_solref',
+          'pair_solreffriction',
+          'qacc',
+          'qacc_smooth',
+          'qacc_warmstart',
+          'qfrc_actuator',
+          'qfrc_applied',
+          'qfrc_bias',
+          'qfrc_constraint',
+          'qfrc_fluid',
+          'qfrc_gravcomp',
+          'qfrc_passive',
+          'qfrc_smooth',
+          'qpos',
+          'qpos0',
+          'qpos_spring',
+          'qvel',
+          'sensordata',
+          'site_pos',
+          'site_quat',
+          'site_xmat',
+          'site_xpos',
+          'subtree_com',
+          'ten_length',
+          'tendon_actfrcrange',
+          'tendon_armature',
+          'tendon_damping',
+          'tendon_dampingpoly',
+          'tendon_frictionloss',
+          'tendon_invweight0',
+          'tendon_length0',
+          'tendon_lengthspring',
+          'tendon_margin',
+          'tendon_range',
+          'tendon_solimp_fri',
+          'tendon_solimp_lim',
+          'tendon_solref_fri',
+          'tendon_solref_lim',
+          'tendon_stiffness',
+          'tendon_stiffnesspoly',
+          'time',
+          'xanchor',
+          'xaxis',
+          'xfrc_applied',
+          'ximat',
+          'xipos',
+          'xmat',
+          'xpos',
+          'xquat',
+      ]),
+      stage_out_argnames=set([
+          'act',
+          'act_dot',
+          'actuator_force',
+          'actuator_length',
+          'cam_xmat',
+          'cam_xpos',
+          'cdof',
+          'cdof_dot',
+          'cvel',
+          'geom_xmat',
+          'geom_xpos',
+          'history',
+          'qacc',
+          'qacc_smooth',
+          'qacc_warmstart',
+          'qfrc_actuator',
+          'qfrc_bias',
+          'qfrc_constraint',
+          'qfrc_fluid',
+          'qfrc_gravcomp',
+          'qfrc_passive',
+          'qfrc_smooth',
+          'qpos',
+          'qvel',
+          'sensordata',
+          'site_xmat',
+          'site_xpos',
+          'subtree_com',
+          'ten_length',
+          'time',
+          'xanchor',
+          'xaxis',
+          'ximat',
+          'xipos',
+          'xmat',
+          'xpos',
+          'xquat',
+      ]),
+      graph_mode=m.opt._impl.graph_mode,
+      has_side_effect=False,
+  )
+  out = jf(
+      d.qpos.shape[0],
+      m._impl.D_colind,
+      m._impl.D_diag,
+      m._impl.D_rowadr,
+      m._impl.D_rownnz,
+      m._impl.M_elemid,
+      m._impl.M_fullm_i,
+      m._impl.M_fullm_j,
+      m._impl.M_fullm_upper_elemid,
+      m._impl.M_fullm_upper_i,
+      m._impl.M_fullm_upper_j,
+      m._impl.M_mulm_col,
+      m._impl.M_mulm_madr,
+      m._impl.M_mulm_rowadr,
+      m._impl.M_rowadr,
+      m._impl.M_rownnz,
+      m._impl.M_tiles,
+      m.actuator_acc0,
+      m.actuator_actadr,
+      m.actuator_actearly,
+      m.actuator_actlimited,
+      m.actuator_actnum,
+      m.actuator_actrange,
+      m.actuator_biasprm,
+      m.actuator_biastype,
+      m.actuator_cranklength,
+      m.actuator_ctrllimited,
+      m.actuator_ctrlrange,
+      m._impl.actuator_delay,
+      m.actuator_dynprm,
+      m.actuator_dyntype,
+      m.actuator_forcelimited,
+      m.actuator_forcerange,
+      m.actuator_gainprm,
+      m.actuator_gaintype,
+      m.actuator_gear,
+      m._impl.actuator_history,
+      m._impl.actuator_historyadr,
+      m.actuator_lengthrange,
+      m.actuator_trnid,
+      m.actuator_trntype,
+      m._impl.actuator_trntype_body_adr,
+      m._impl.block_dim,
+      m._impl.body_branch_start,
+      m._impl.body_branches,
+      m.body_dofadr,
+      m.body_dofnum,
+      m._impl.body_fluid_ellipsoid,
+      m.body_geomadr,
+      m.body_geomnum,
+      m.body_gravcomp,
+      m.body_inertia,
+      m.body_invweight0,
+      m.body_ipos,
+      m.body_iquat,
+      m._impl.body_isdofancestor,
+      m.body_jntadr,
+      m.body_jntnum,
+      m.body_mass,
+      m.body_mocapid,
+      m.body_parentid,
+      m.body_pos,
+      m.body_quat,
+      m.body_rootid,
+      m.body_subtreemass,
+      m._impl.body_tree,
+      m.body_treeid,
+      m.body_weldid,
+      m.cam_bodyid,
+      m.cam_fovy,
+      m.cam_intrinsic,
+      m.cam_mat0,
+      m.cam_mode,
+      m.cam_pos,
+      m.cam_pos0,
+      m.cam_poscom0,
+      m.cam_quat,
+      m.cam_resolution,
+      m.cam_sensorsize,
+      m.cam_targetbodyid,
+      m.dof_armature,
+      m.dof_bodyid,
+      m.dof_damping,
+      m.dof_dampingpoly,
+      m.dof_frictionloss,
+      m.dof_invweight0,
+      m.dof_jntid,
+      m.dof_parentid,
+      m.dof_solimp,
+      m.dof_solref,
+      m.dof_treeid,
+      m._impl.dof_tri_col,
+      m._impl.dof_tri_row,
+      m._impl.eq_connect_adr,
+      m.eq_data,
+      m._impl.eq_flex_adr,
+      m._impl.eq_jnt_adr,
+      m.eq_obj1id,
+      m.eq_obj2id,
+      m.eq_objtype,
+      m.eq_solimp,
+      m.eq_solref,
+      m._impl.eq_ten_adr,
+      m.eq_type,
+      m._impl.eq_wld_adr,
+      m._impl.flex_bending,
+      m._impl.flex_bendingadr,
+      m._impl.flex_centered,
+      m._impl.flex_conaffinity,
+      m._impl.flex_condim,
+      m._impl.flex_contype,
+      m._impl.flex_damping,
+      m._impl.flex_dim,
+      m._impl.flex_edge,
+      m._impl.flex_edgeadr,
+      m._impl.flex_edgeflap,
+      m._impl.flex_edgenum,
+      m._impl.flex_elem,
+      m._impl.flex_elemadr,
+      m._impl.flex_elemdataadr,
+      m._impl.flex_elemedge,
+      m._impl.flex_elemedgeadr,
+      m._impl.flex_elemnum,
+      m._impl.flex_friction,
+      m._impl.flex_gap,
+      m._impl.flex_margin,
+      m._impl.flex_priority,
+      m._impl.flex_radius,
+      m._impl.flex_shell,
+      m._impl.flex_shelldataadr,
+      m._impl.flex_shellnum,
+      m._impl.flex_solimp,
+      m._impl.flex_solmix,
+      m._impl.flex_solref,
+      m._impl.flex_stiffness,
+      m._impl.flex_stiffnessadr,
+      m._impl.flex_vert,
+      m.flex_vertadr,
+      m._impl.flex_vertbodyid,
+      m._impl.flex_vertflexid,
+      m.flex_vertnum,
+      m._impl.flexedge_J_colind,
+      m._impl.flexedge_J_rowadr,
+      m._impl.flexedge_J_rownnz,
+      m._impl.flexedge_invweight0,
+      m._impl.flexedge_length0,
+      m.geom_aabb,
+      m.geom_bodyid,
+      m.geom_conaffinity,
+      m.geom_condim,
+      m.geom_contype,
+      jax.numpy.expand_dims(m.geom_dataid, 0),
+      m.geom_fluid,
+      m.geom_friction,
+      m.geom_gap,
+      m.geom_group,
+      m.geom_margin,
+      m.geom_matid,
+      m._impl.geom_pair_type_count,
+      m._impl.geom_plugin_index,
+      m.geom_pos,
+      m.geom_priority,
+      m.geom_quat,
+      m.geom_rbound,
+      m.geom_rgba,
+      m.geom_size,
+      m.geom_solimp,
+      m.geom_solmix,
+      m.geom_solref,
+      m.geom_type,
+      m._impl.has_fluid,
+      m._impl.has_sdf_geom,
+      m.hfield_adr,
+      m.hfield_data,
+      m.hfield_ncol,
+      m.hfield_nrow,
+      m.hfield_size,
+      m._impl.is_sparse,
+      m.jnt_actfrclimited,
+      m.jnt_actfrcrange,
+      m.jnt_actgravcomp,
+      m.jnt_axis,
+      m.jnt_bodyid,
+      m.jnt_dofadr,
+      m._impl.jnt_limited_ball_adr,
+      m._impl.jnt_limited_slide_hinge_adr,
+      m.jnt_margin,
+      m.jnt_pos,
+      m.jnt_qposadr,
+      m.jnt_range,
+      m.jnt_solimp,
+      m.jnt_solref,
+      m.jnt_stiffness,
+      m.jnt_stiffnesspoly,
+      m.jnt_type,
+      m._impl.light_bodyid,
+      m.light_dir,
+      m.light_dir0,
+      m.light_mode,
+      m.light_pos,
+      m.light_pos0,
+      m.light_poscom0,
+      m._impl.light_targetbodyid,
+      m._impl.mapM2D,
+      m.mat_rgba,
+      m._impl.max_ten_J_rownnz,
+      m.mesh_face,
+      m.mesh_faceadr,
+      m.mesh_graph,
+      m.mesh_graphadr,
+      m.mesh_normal,
+      m.mesh_normaladr,
+      m.mesh_normalnum,
+      m.mesh_octadr,
+      m._impl.mesh_polyadr,
+      m._impl.mesh_polymap,
+      m._impl.mesh_polymapadr,
+      m._impl.mesh_polymapnum,
+      m._impl.mesh_polynormal,
+      m._impl.mesh_polynum,
+      m._impl.mesh_polyvert,
+      m._impl.mesh_polyvertadr,
+      m._impl.mesh_polyvertnum,
+      m.mesh_quat,
+      m.mesh_vert,
+      m.mesh_vertadr,
+      m.mesh_vertnum,
+      m.nC,
+      m.nD,
+      m.nJten,
+      m.na,
+      m._impl.nacttrnbody,
+      m.nbody,
+      m._impl.nbranch,
+      m.ncam,
+      m.neq,
+      m.nflex,
+      m._impl.nflexedge,
+      m._impl.nflexelem,
+      m._impl.nflexshelldata,
+      m._impl.nflexvert,
+      m.ngeom,
+      m.ngravcomp,
+      m.nhistory,
+      m.njnt,
+      m.nlight,
+      m._impl.nmaxcondim,
+      m._impl.nmaxmeshdeg,
+      m._impl.nmaxpolygon,
+      m._impl.nmaxpyramid,
+      m.nmeshface,
+      m._impl.nrangefinder,
+      m._impl.nsensorcollision,
+      m._impl.nsensorcontact,
+      m._impl.nsensortaxel,
+      m.nsite,
+      m.ntendon,
+      m._impl.ntree,
+      m.nu,
+      m.nv,
+      m._impl.nv_pad,
+      m.nwrap,
+      m._impl.nxn_geom_pair_filtered,
+      m._impl.nxn_pairid,
+      m._impl.nxn_pairid_filtered,
+      m._impl.oct_aabb,
+      m._impl.oct_child,
+      m._impl.oct_coeff,
+      m.pair_dim,
+      m.pair_friction,
+      m.pair_gap,
+      m.pair_margin,
+      m.pair_solimp,
+      m.pair_solref,
+      m.pair_solreffriction,
+      m._impl.plugin,
+      m._impl.plugin_attr,
+      m._impl.qD_fullm_i,
+      m._impl.qD_fullm_j,
+      m._impl.qLD_all_updates,
+      m._impl.qLD_level_offsets,
+      m._impl.qLD_updates,
+      m.qpos0,
+      m.qpos_spring,
+      m._impl.rangefinder_sensor_adr,
+      m._impl.sensor_acc_adr,
+      m.sensor_adr,
+      m._impl.sensor_adr_to_contact_adr,
+      m._impl.sensor_contact_adr,
+      m.sensor_cutoff,
+      m.sensor_datatype,
+      m._impl.sensor_delay,
+      m.sensor_dim,
+      m._impl.sensor_e_kinetic,
+      m._impl.sensor_e_potential,
+      m._impl.sensor_history,
+      m._impl.sensor_historyadr,
+      m._impl.sensor_interval,
+      m.sensor_intprm,
+      m._impl.sensor_limitfrc_adr,
+      m._impl.sensor_limitpos_adr,
+      m._impl.sensor_limitvel_adr,
+      m.sensor_objid,
+      m.sensor_objtype,
+      m._impl.sensor_pos_adr,
+      m._impl.sensor_rangefinder_adr,
+      m._impl.sensor_rangefinder_bodyid,
+      m.sensor_refid,
+      m.sensor_reftype,
+      m._impl.sensor_rne_postconstraint,
+      m._impl.sensor_subtree_vel,
+      m._impl.sensor_tendonactfrc_adr,
+      m._impl.sensor_touch_adr,
+      m.sensor_type,
+      m._impl.sensor_vel_adr,
+      m.site_bodyid,
+      m.site_pos,
+      m.site_quat,
+      m.site_size,
+      m.site_type,
+      m._impl.taxel_sensorid,
+      m._impl.taxel_vertadr,
+      m._impl.ten_J_colind,
+      m._impl.ten_J_rowadr,
+      m._impl.ten_J_rownnz,
+      m.tendon_actfrclimited,
+      m.tendon_actfrcrange,
+      m.tendon_adr,
+      m.tendon_armature,
+      m.tendon_damping,
+      m.tendon_dampingpoly,
+      m.tendon_frictionloss,
+      m._impl.tendon_geom_adr,
+      m.tendon_invweight0,
+      m._impl.tendon_jnt_adr,
+      m.tendon_length0,
+      m.tendon_lengthspring,
+      m._impl.tendon_limited_adr,
+      m.tendon_margin,
+      m.tendon_num,
+      m.tendon_range,
+      m._impl.tendon_site_pair_adr,
+      m.tendon_solimp_fri,
+      m.tendon_solimp_lim,
+      m.tendon_solref_fri,
+      m.tendon_solref_lim,
+      m.tendon_stiffness,
+      m.tendon_stiffnesspoly,
+      m._impl.wrap_geom_adr,
+      m._impl.wrap_jnt_adr,
+      m.wrap_objid,
+      m.wrap_prm,
+      m._impl.wrap_pulley_scale,
+      m._impl.wrap_site_pair_adr,
+      m.wrap_type,
+      m.opt._impl.broadphase,
+      m.opt._impl.broadphase_filter,
+      m.opt._impl.ccd_iterations,
+      m.opt._impl.ccd_tolerance,
+      m.opt.cone,
+      m.opt._impl.contact_sensor_maxmatch,
+      m.opt.density,
+      m.opt.disableflags,
+      m.opt.enableflags,
+      m.opt._impl.graph_conditional,
+      m.opt.gravity,
+      m.opt._impl.impratio_invsqrt,
+      m.opt.integrator,
+      m.opt.iterations,
+      m.opt.ls_iterations,
+      m.opt._impl.ls_parallel,
+      m.opt._impl.ls_parallel_min_step,
+      m.opt.ls_tolerance,
+      m.opt.magnetic,
+      m.opt._impl.run_collision_detection,
+      m.opt._impl.sdf_initpoints,
+      m.opt._impl.sdf_iterations,
+      m.opt.solver,
+      m.opt.timestep,
+      m.opt.tolerance,
+      m.opt.viscosity,
+      m.opt.wind,
+      m.stat.meaninertia,
+      d._impl.naccdmax,
+      d._impl.naconmax,
+      d._impl.njmax,
+      d._impl.njmax_nnz,
+      d._impl.M,
+      d.act,
+      d.act_dot,
+      d.actuator_force,
+      d.actuator_length,
+      d._impl.actuator_moment,
+      d._impl.actuator_velocity,
+      d._impl.cacc,
+      d.cam_xmat,
+      d.cam_xpos,
+      d.cdof,
+      d.cdof_dot,
+      d._impl.cfrc_ext,
+      d._impl.cfrc_int,
+      d._impl.cinert,
+      d._impl.crb,
+      d.ctrl,
+      d.cvel,
+      d._impl.dof_island,
+      d._impl.dof_islandid,
+      d._impl.efc_islandid,
+      d._impl.energy,
+      d.eq_active,
+      d._impl.flexedge_J,
+      d._impl.flexedge_length,
+      d._impl.flexedge_velocity,
+      d._impl.flexvert_xpos,
+      d.geom_xmat,
+      d.geom_xpos,
+      d.history,
+      d._impl.iqacc,
+      d._impl.iqacc_smooth,
+      d._impl.iqfrc_constraint,
+      d._impl.iqfrc_smooth,
+      d._impl.island_dofadr,
+      d._impl.island_efcadr,
+      d._impl.island_ne,
+      d._impl.island_nefc,
+      d._impl.island_nf,
+      d._impl.island_nv,
+      d._impl.light_xdir,
+      d._impl.light_xpos,
+      d._impl.map_dof2idof,
+      d._impl.map_efc2iefc,
+      d._impl.map_idof2dof,
+      d._impl.map_iefc2efc,
+      d.mocap_pos,
+      d.mocap_quat,
+      d._impl.moment_colind,
+      d._impl.moment_rowadr,
+      d._impl.moment_rownnz,
+      d._impl.nacon,
+      d._impl.ncollision,
+      d._impl.ne,
+      d._impl.nefc,
+      d._impl.nf,
+      d._impl.nidof,
+      d._impl.nisland,
+      d._impl.nl,
+      d._impl.qLD,
+      d._impl.qLDiagInv,
+      d._impl.qLU,
+      d.qacc,
+      d.qacc_smooth,
+      d.qacc_warmstart,
+      d.qfrc_actuator,
+      d.qfrc_applied,
+      d.qfrc_bias,
+      d.qfrc_constraint,
+      d._impl.qfrc_damper,
+      d.qfrc_fluid,
+      d.qfrc_gravcomp,
+      d.qfrc_passive,
+      d.qfrc_smooth,
+      d._impl.qfrc_spring,
+      d.qpos,
+      d.qvel,
+      d.sensordata,
+      d.site_xmat,
+      d.site_xpos,
+      d._impl.solver_niter,
+      d._impl.subtree_angmom,
+      d.subtree_com,
+      d._impl.subtree_linvel,
+      d._impl.ten_J,
+      d.ten_length,
+      d._impl.ten_velocity,
+      d._impl.ten_wrapadr,
+      d._impl.ten_wrapnum,
+      d.time,
+      d._impl.tree_island,
+      d._impl.wrap_obj,
+      d._impl.wrap_xpos,
+      d.xanchor,
+      d.xaxis,
+      d.xfrc_applied,
+      d.ximat,
+      d.xipos,
+      d.xmat,
+      d.xpos,
+      d.xquat,
+      d._impl.contact__dim,
+      d._impl.contact__dist,
+      d._impl.contact__efc_address,
+      d._impl.contact__flex,
+      d._impl.contact__frame,
+      d._impl.contact__friction,
+      d._impl.contact__geom,
+      d._impl.contact__geomcollisionid,
+      d._impl.contact__includemargin,
+      d._impl.contact__pos,
+      d._impl.contact__solimp,
+      d._impl.contact__solref,
+      d._impl.contact__solreffriction,
+      d._impl.contact__type,
+      d._impl.contact__vert,
+      d._impl.contact__worldid,
+      d._impl.efc__D,
+      d._impl.efc__J,
+      d._impl.efc__J_colind,
+      d._impl.efc__J_rowadr,
+      d._impl.efc__J_rownnz,
+      d._impl.efc__Jqvel,
+      d._impl.efc__Ma,
+      d._impl.efc__aref,
+      d._impl.efc__force,
+      d._impl.efc__frictionloss,
+      d._impl.efc__iD,
+      d._impl.efc__iJ,
+      d._impl.efc__iJ_colind,
+      d._impl.efc__iJ_rowadr,
+      d._impl.efc__iJ_rownnz,
+      d._impl.efc__iaref,
+      d._impl.efc__id,
+      d._impl.efc__iforce,
+      d._impl.efc__ifrictionloss,
+      d._impl.efc__iid,
+      d._impl.efc__island,
+      d._impl.efc__istate,
+      d._impl.efc__itype,
+      d._impl.efc__margin,
+      d._impl.efc__pos,
+      d._impl.efc__state,
+      d._impl.efc__type,
+      d._impl.efc__vel,
+  )
+  d = d.tree_replace({
+      '_impl.M': out[0],
+      'act': out[1],
+      'act_dot': out[2],
+      'actuator_force': out[3],
+      'actuator_length': out[4],
+      '_impl.actuator_moment': out[5],
+      '_impl.actuator_velocity': out[6],
+      '_impl.cacc': out[7],
+      'cam_xmat': out[8],
+      'cam_xpos': out[9],
+      'cdof': out[10],
+      'cdof_dot': out[11],
+      '_impl.cfrc_ext': out[12],
+      '_impl.cfrc_int': out[13],
+      '_impl.cinert': out[14],
+      '_impl.crb': out[15],
+      'cvel': out[16],
+      '_impl.dof_island': out[17],
+      '_impl.dof_islandid': out[18],
+      '_impl.efc_islandid': out[19],
+      '_impl.energy': out[20],
+      '_impl.flexedge_J': out[21],
+      '_impl.flexedge_length': out[22],
+      '_impl.flexedge_velocity': out[23],
+      '_impl.flexvert_xpos': out[24],
+      'geom_xmat': out[25],
+      'geom_xpos': out[26],
+      'history': out[27],
+      '_impl.iqacc': out[28],
+      '_impl.iqacc_smooth': out[29],
+      '_impl.iqfrc_constraint': out[30],
+      '_impl.iqfrc_smooth': out[31],
+      '_impl.island_dofadr': out[32],
+      '_impl.island_efcadr': out[33],
+      '_impl.island_ne': out[34],
+      '_impl.island_nefc': out[35],
+      '_impl.island_nf': out[36],
+      '_impl.island_nv': out[37],
+      '_impl.light_xdir': out[38],
+      '_impl.light_xpos': out[39],
+      '_impl.map_dof2idof': out[40],
+      '_impl.map_efc2iefc': out[41],
+      '_impl.map_idof2dof': out[42],
+      '_impl.map_iefc2efc': out[43],
+      '_impl.moment_colind': out[44],
+      '_impl.moment_rowadr': out[45],
+      '_impl.moment_rownnz': out[46],
+      '_impl.nacon': out[47],
+      '_impl.ncollision': out[48],
+      '_impl.ne': out[49],
+      '_impl.nefc': out[50],
+      '_impl.nf': out[51],
+      '_impl.nidof': out[52],
+      '_impl.nisland': out[53],
+      '_impl.nl': out[54],
+      '_impl.qLD': out[55],
+      '_impl.qLDiagInv': out[56],
+      '_impl.qLU': out[57],
+      'qacc': out[58],
+      'qacc_smooth': out[59],
+      'qacc_warmstart': out[60],
+      'qfrc_actuator': out[61],
+      'qfrc_bias': out[62],
+      'qfrc_constraint': out[63],
+      '_impl.qfrc_damper': out[64],
+      'qfrc_fluid': out[65],
+      'qfrc_gravcomp': out[66],
+      'qfrc_passive': out[67],
+      'qfrc_smooth': out[68],
+      '_impl.qfrc_spring': out[69],
+      'qpos': out[70],
+      'qvel': out[71],
+      'sensordata': out[72],
+      'site_xmat': out[73],
+      'site_xpos': out[74],
+      '_impl.solver_niter': out[75],
+      '_impl.subtree_angmom': out[76],
+      'subtree_com': out[77],
+      '_impl.subtree_linvel': out[78],
+      '_impl.ten_J': out[79],
+      'ten_length': out[80],
+      '_impl.ten_velocity': out[81],
+      '_impl.ten_wrapadr': out[82],
+      '_impl.ten_wrapnum': out[83],
+      'time': out[84],
+      '_impl.tree_island': out[85],
+      '_impl.wrap_obj': out[86],
+      '_impl.wrap_xpos': out[87],
+      'xanchor': out[88],
+      'xaxis': out[89],
+      'ximat': out[90],
+      'xipos': out[91],
+      'xmat': out[92],
+      'xpos': out[93],
+      'xquat': out[94],
+      '_impl.contact__dim': out[95],
+      '_impl.contact__dist': out[96],
+      '_impl.contact__efc_address': out[97],
+      '_impl.contact__flex': out[98],
+      '_impl.contact__frame': out[99],
+      '_impl.contact__friction': out[100],
+      '_impl.contact__geom': out[101],
+      '_impl.contact__geomcollisionid': out[102],
+      '_impl.contact__includemargin': out[103],
+      '_impl.contact__pos': out[104],
+      '_impl.contact__solimp': out[105],
+      '_impl.contact__solref': out[106],
+      '_impl.contact__solreffriction': out[107],
+      '_impl.contact__type': out[108],
+      '_impl.contact__vert': out[109],
+      '_impl.contact__worldid': out[110],
+      '_impl.efc__D': out[111],
+      '_impl.efc__J': out[112],
+      '_impl.efc__J_colind': out[113],
+      '_impl.efc__J_rowadr': out[114],
+      '_impl.efc__J_rownnz': out[115],
+      '_impl.efc__Jqvel': out[116],
+      '_impl.efc__Ma': out[117],
+      '_impl.efc__aref': out[118],
+      '_impl.efc__force': out[119],
+      '_impl.efc__frictionloss': out[120],
+      '_impl.efc__iD': out[121],
+      '_impl.efc__iJ': out[122],
+      '_impl.efc__iJ_colind': out[123],
+      '_impl.efc__iJ_rowadr': out[124],
+      '_impl.efc__iJ_rownnz': out[125],
+      '_impl.efc__iaref': out[126],
+      '_impl.efc__id': out[127],
+      '_impl.efc__iforce': out[128],
+      '_impl.efc__ifrictionloss': out[129],
+      '_impl.efc__iid': out[130],
+      '_impl.efc__island': out[131],
+      '_impl.efc__istate': out[132],
+      '_impl.efc__itype': out[133],
+      '_impl.efc__margin': out[134],
+      '_impl.efc__pos': out[135],
+      '_impl.efc__state': out[136],
+      '_impl.efc__type': out[137],
+      '_impl.efc__vel': out[138],
   })
   return d
 
diff --git a/mjx/mujoco/mjx/warp/forward_test.py b/mjx/mujoco/mjx/warp/forward_test.py
index 2c38891e..95d49e0f 100644
--- a/mjx/mujoco/mjx/warp/forward_test.py
+++ b/mjx/mujoco/mjx/warp/forward_test.py
@@ -176,8 +176,8 @@ class ForwardTest(parameterized.TestCase):
 
       qm = np.zeros((m.nv, m.nv))
       mujoco.mju_sym2dense(qm, d.M, m.M_rownnz, m.M_rowadr, m.M_colind)
-      # mjwarp adds padding to qM
-      tu.assert_eq(qm, dx._impl.qM[: m.nv, : m.nv], 'qM')
+      # mjwarp adds padding to M
+      tu.assert_eq(qm, dx._impl.M[: m.nv, : m.nv], 'M')
       # qLD is fused in a cholesky factorize and solve, and not written to.
 
       tu.assert_contact_eq(d, dx, worldid=i)
diff --git a/mjx/mujoco/mjx/warp/smooth_test.py b/mjx/mujoco/mjx/warp/smooth_test.py
index 7c8bb368..442461cb 100644
--- a/mjx/mujoco/mjx/warp/smooth_test.py
+++ b/mjx/mujoco/mjx/warp/smooth_test.py
@@ -112,7 +112,7 @@ class SmoothTest(parameterized.TestCase):
 
     d = mujoco.MjData(m)
     mx = mjx.put_model(m, impl='warp')
-    mx = mx.replace(_impl=mx._impl.replace(qM_tiles=()))
+    mx = mx.replace(_impl=mx._impl.replace(M_tiles=()))
 
     rng = jax.random.PRNGKey(0)
     dx = mjx.make_data(m, impl='warp')
@@ -188,7 +188,7 @@ class SmoothTest(parameterized.TestCase):
     batch_size = 7
     d = mujoco.MjData(m)
     mx = mjx.put_model(m, impl='warp')
-    mx = mx.replace(_impl=mx._impl.replace(qM_tiles=()))
+    mx = mx.replace(_impl=mx._impl.replace(M_tiles=()))
 
     worldids = jp.arange(batch_size)
     dx_batch = jax.vmap(functools.partial(tu.make_data, m))(worldids)
diff --git a/mjx/mujoco/mjx/warp/types.py b/mjx/mujoco/mjx/warp/types.py
index df4098d2..5e7ef094 100644
--- a/mjx/mujoco/mjx/warp/types.py
+++ b/mjx/mujoco/mjx/warp/types.py
@@ -78,6 +78,29 @@ class BlockDim:
   """Block dimension 'block_dim' settings for wp.launch_tiled.
 
   TODO(team): experimental and may be removed
+
+  Attributes:
+    segmented_sort: segmented sort block dimension (collision_driver)
+    euler_dense: Euler dense block dimension (forward)
+    actuator_velocity: actuator velocity block dimension (forward)
+    ray: ray block dimension (ray)
+    contact_sort: contact sort block dimension (sensor)
+    energy_vel_kinetic: energy velocity kinetic block dimension (sensor)
+    cholesky_factorize: Cholesky factorize block dimension (smooth)
+    cholesky_solve: Cholesky solve block dimension (smooth)
+    cholesky_factorize_solve: Cholesky factorize and solve block dimension
+      (smooth)
+    solve_LD_sparse_fused: solve LD sparse fused block dimension (smooth)
+    update_gradient_cholesky: update gradient Cholesky block dimension (solver)
+    update_gradient_cholesky_blocked: update gradient Cholesky blocked block
+      dimension (solver)
+    update_gradient_JTDAJ_sparse: update gradient JTDAJ sparse block dimension
+      (solver)
+    update_gradient_JTDAJ_dense: update gradient JTDAJ dense block dimension
+      (solver)
+    linesearch_iterative: linesearch iterative block dimension (solver)
+    contact_jac_tiled: contact Jacobian tiled block dimension (solver)
+    qderiv_actuator_dense: qderiv actuator dense block dimension (derivative)
   """
 
   actuator_velocity: int
@@ -134,10 +157,26 @@ class OptionWarp(PyTreeNode):
 
 class ModelWarp(PyTreeNode):
   """Derived fields from Model."""
-
+  D_colind: np.ndarray
+  D_diag: np.ndarray
+  D_rowadr: np.ndarray
+  D_rownnz: np.ndarray
   M_colind: np.ndarray
+  M_elemid: np.ndarray
+  M_fullm_i: np.ndarray
+  M_fullm_j: np.ndarray
+  M_fullm_upper_elemid: np.ndarray
+  M_fullm_upper_i: np.ndarray
+  M_fullm_upper_j: np.ndarray
+  M_mulm_col: np.ndarray
+  M_mulm_madr: np.ndarray
+  M_mulm_rowadr: np.ndarray
   M_rowadr: np.ndarray
   M_rownnz: np.ndarray
+  M_tiles: Tuple[TileSet, ...]
+  actuator_delay: np.ndarray
+  actuator_history: np.ndarray
+  actuator_historyadr: np.ndarray
   actuator_trntype_body_adr: np.ndarray
   block_dim: BlockDim
   body_branch_start: np.ndarray
@@ -204,6 +243,8 @@ class ModelWarp(PyTreeNode):
   light_active: jax.Array
   light_bodyid: np.ndarray
   light_targetbodyid: np.ndarray
+  mapD2M: np.ndarray
+  mapM2D: np.ndarray
   mapM2M: np.ndarray
   mat_texrepeat: jax.Array
   max_ten_J_rownnz: int
@@ -249,22 +290,22 @@ class ModelWarp(PyTreeNode):
   oct_coeff: np.ndarray
   plugin: np.ndarray
   plugin_attr: np.ndarray
+  qD_fullm_i: np.ndarray
+  qD_fullm_j: np.ndarray
   qLD_all_updates: np.ndarray
   qLD_level_offsets: np.ndarray
   qLD_updates: Tuple[np.ndarray, ...]
-  qM_fullm_i: np.ndarray
-  qM_fullm_j: np.ndarray
-  qM_mulm_col: np.ndarray
-  qM_mulm_madr: np.ndarray
-  qM_mulm_rowadr: np.ndarray
-  qM_tiles: Tuple[TileSet, ...]
   rangefinder_sensor_adr: np.ndarray
   sensor_acc_adr: np.ndarray
   sensor_adr_to_contact_adr: np.ndarray
   sensor_collision_start_adr: np.ndarray
   sensor_contact_adr: np.ndarray
+  sensor_delay: np.ndarray
   sensor_e_kinetic: bool
   sensor_e_potential: bool
+  sensor_history: np.ndarray
+  sensor_historyadr: np.ndarray
+  sensor_interval: np.ndarray
   sensor_limitfrc_adr: np.ndarray
   sensor_limitpos_adr: np.ndarray
   sensor_limitvel_adr: np.ndarray
@@ -299,7 +340,7 @@ class ModelWarp(PyTreeNode):
 
 class DataWarp(PyTreeNode):
   """Derived fields from Data."""
-
+  M: jax.Array
   actuator_moment: jax.Array
   actuator_velocity: jax.Array
   cacc: jax.Array
@@ -323,6 +364,8 @@ class DataWarp(PyTreeNode):
   contact__vert: jax.Array
   contact__worldid: jax.Array
   crb: jax.Array
+  dof_island: jax.Array
+  dof_islandid: jax.Array
   efc__D: jax.Array
   efc__J: jax.Array
   efc__J_colind: jax.Array
@@ -333,19 +376,46 @@ class DataWarp(PyTreeNode):
   efc__aref: jax.Array
   efc__force: jax.Array
   efc__frictionloss: jax.Array
+  efc__iD: jax.Array
+  efc__iJ: jax.Array
+  efc__iJ_colind: jax.Array
+  efc__iJ_rowadr: jax.Array
+  efc__iJ_rownnz: jax.Array
+  efc__iaref: jax.Array
   efc__id: jax.Array
+  efc__iforce: jax.Array
+  efc__ifrictionloss: jax.Array
+  efc__iid: jax.Array
+  efc__island: jax.Array
+  efc__istate: jax.Array
+  efc__itype: jax.Array
   efc__margin: jax.Array
   efc__pos: jax.Array
   efc__state: jax.Array
   efc__type: jax.Array
   efc__vel: jax.Array
+  efc_islandid: jax.Array
   energy: jax.Array
   flexedge_J: jax.Array
   flexedge_length: jax.Array
   flexedge_velocity: jax.Array
   flexvert_xpos: jax.Array
+  iqacc: jax.Array
+  iqacc_smooth: jax.Array
+  iqfrc_constraint: jax.Array
+  iqfrc_smooth: jax.Array
+  island_dofadr: jax.Array
+  island_efcadr: jax.Array
+  island_ne: jax.Array
+  island_nefc: jax.Array
+  island_nf: jax.Array
+  island_nv: jax.Array
   light_xdir: jax.Array
   light_xpos: jax.Array
+  map_dof2idof: jax.Array
+  map_efc2iefc: jax.Array
+  map_idof2dof: jax.Array
+  map_iefc2efc: jax.Array
   moment_colind: jax.Array
   moment_rowadr: jax.Array
   moment_rownnz: jax.Array
@@ -356,6 +426,7 @@ class DataWarp(PyTreeNode):
   ne: jax.Array
   nefc: jax.Array
   nf: jax.Array
+  nidof: jax.Array
   nisland: jax.Array
   njmax: int
   njmax_nnz: int
@@ -364,7 +435,7 @@ class DataWarp(PyTreeNode):
   nworld: int
   qLD: jax.Array
   qLDiagInv: jax.Array
-  qM: jax.Array
+  qLU: jax.Array
   qfrc_damper: jax.Array
   qfrc_spring: jax.Array
   solver_niter: jax.Array
@@ -401,6 +472,7 @@ DATA_NON_VMAP = {
     'nacon',
     'naconmax',
     'ncollision',
+    'nidof',
     'njmax',
     'njmax_nnz',
     'njmax_pad',
@@ -434,6 +506,7 @@ batching.register_vmappable(DataWarp, int, int, _to_elt, _from_elt, None)
 
 _NDIM = {
     'Data': {
+        'M': 3,
         'act': 2,
         'act_dot': 2,
         'actuator_force': 2,
@@ -467,6 +540,8 @@ _NDIM = {
         'crb': 3,
         'ctrl': 2,
         'cvel': 3,
+        'dof_island': 2,
+        'dof_islandid': 2,
         'efc__D': 2,
         'efc__J': 3,
         'efc__J_colind': 3,
@@ -477,12 +552,25 @@ _NDIM = {
         'efc__aref': 2,
         'efc__force': 2,
         'efc__frictionloss': 2,
+        'efc__iD': 2,
+        'efc__iJ': 3,
+        'efc__iJ_colind': 3,
+        'efc__iJ_rowadr': 2,
+        'efc__iJ_rownnz': 2,
+        'efc__iaref': 2,
         'efc__id': 2,
+        'efc__iforce': 2,
+        'efc__ifrictionloss': 2,
+        'efc__iid': 2,
+        'efc__island': 2,
+        'efc__istate': 2,
+        'efc__itype': 2,
         'efc__margin': 2,
         'efc__pos': 2,
         'efc__state': 2,
         'efc__type': 2,
         'efc__vel': 2,
+        'efc_islandid': 2,
         'energy': 2,
         'eq_active': 2,
         'flexedge_J': 2,
@@ -491,8 +579,23 @@ _NDIM = {
         'flexvert_xpos': 3,
         'geom_xmat': 4,
         'geom_xpos': 3,
+        'history': 2,
+        'iqacc': 2,
+        'iqacc_smooth': 2,
+        'iqfrc_constraint': 2,
+        'iqfrc_smooth': 2,
+        'island_dofadr': 2,
+        'island_efcadr': 2,
+        'island_ne': 2,
+        'island_nefc': 2,
+        'island_nf': 2,
+        'island_nv': 2,
         'light_xdir': 3,
         'light_xpos': 3,
+        'map_dof2idof': 2,
+        'map_efc2iefc': 2,
+        'map_idof2dof': 2,
+        'map_iefc2efc': 2,
         'mocap_pos': 3,
         'mocap_quat': 3,
         'moment_colind': 2,
@@ -505,6 +608,7 @@ _NDIM = {
         'ne': 1,
         'nefc': 1,
         'nf': 1,
+        'nidof': 1,
         'nisland': 1,
         'njmax': 0,
         'njmax_nnz': 0,
@@ -513,7 +617,7 @@ _NDIM = {
         'nworld': 0,
         'qLD': 3,
         'qLDiagInv': 2,
-        'qM': 3,
+        'qLU': 3,
         'qacc': 2,
         'qacc_smooth': 2,
         'qacc_warmstart': 2,
@@ -556,9 +660,23 @@ _NDIM = {
         'xquat': 3,
     },
     'Model': {
+        'D_colind': 1,
+        'D_diag': 1,
+        'D_rowadr': 1,
+        'D_rownnz': 1,
         'M_colind': 1,
+        'M_elemid': 2,
+        'M_fullm_i': 1,
+        'M_fullm_j': 1,
+        'M_fullm_upper_elemid': 1,
+        'M_fullm_upper_i': 1,
+        'M_fullm_upper_j': 1,
+        'M_mulm_col': 1,
+        'M_mulm_madr': 1,
+        'M_mulm_rowadr': 1,
         'M_rowadr': 1,
         'M_rownnz': 1,
+        'M_tiles': -1,
         'actuator_acc0': 2,
         'actuator_actadr': 1,
         'actuator_actearly': 1,
@@ -570,6 +688,7 @@ _NDIM = {
         'actuator_cranklength': 2,
         'actuator_ctrllimited': 1,
         'actuator_ctrlrange': 3,
+        'actuator_delay': 1,
         'actuator_dynprm': 3,
         'actuator_dyntype': 1,
         'actuator_forcelimited': 1,
@@ -577,6 +696,8 @@ _NDIM = {
         'actuator_gainprm': 3,
         'actuator_gaintype': 1,
         'actuator_gear': 3,
+        'actuator_history': 2,
+        'actuator_historyadr': 1,
         'actuator_lengthrange': 3,
         'actuator_trnid': 2,
         'actuator_trntype': 1,
@@ -769,6 +890,8 @@ _NDIM = {
         'light_poscom0': 3,
         'light_targetbodyid': 1,
         'light_type': 2,
+        'mapD2M': 1,
+        'mapM2D': 1,
         'mapM2M': 1,
         'mat_rgba': 3,
         'mat_texid': 3,
@@ -797,6 +920,7 @@ _NDIM = {
         'mesh_vertnum': 1,
         'mocap_bodyid': 1,
         'nC': 0,
+        'nD': 0,
         'nJfe': 0,
         'nJmom': 0,
         'nJten': 0,
@@ -821,6 +945,7 @@ _NDIM = {
         'ngravcomp': 0,
         'nhfield': 0,
         'nhfielddata': 0,
+        'nhistory': 0,
         'njnt': 0,
         'nlight': 0,
         'nmat': 0,
@@ -897,15 +1022,11 @@ _NDIM = {
         'pair_solreffriction': 3,
         'plugin': 1,
         'plugin_attr': 2,
+        'qD_fullm_i': 1,
+        'qD_fullm_j': 1,
         'qLD_all_updates': 2,
         'qLD_level_offsets': 1,
         'qLD_updates': -1,
-        'qM_fullm_i': 1,
-        'qM_fullm_j': 1,
-        'qM_mulm_col': 1,
-        'qM_mulm_madr': 1,
-        'qM_mulm_rowadr': 1,
-        'qM_tiles': -1,
         'qpos0': 2,
         'qpos_spring': 2,
         'rangefinder_sensor_adr': 1,
@@ -916,9 +1037,13 @@ _NDIM = {
         'sensor_contact_adr': 1,
         'sensor_cutoff': 1,
         'sensor_datatype': 1,
+        'sensor_delay': 1,
         'sensor_dim': 1,
         'sensor_e_kinetic': 0,
         'sensor_e_potential': 0,
+        'sensor_history': 2,
+        'sensor_historyadr': 1,
+        'sensor_interval': 2,
         'sensor_intprm': 2,
         'sensor_limitfrc_adr': 1,
         'sensor_limitpos_adr': 1,
@@ -1016,6 +1141,7 @@ _NDIM = {
 }
 _BATCH_DIM = {
     'Data': {
+        'M': True,
         'act': True,
         'act_dot': True,
         'actuator_force': True,
@@ -1049,6 +1175,8 @@ _BATCH_DIM = {
         'crb': True,
         'ctrl': True,
         'cvel': True,
+        'dof_island': True,
+        'dof_islandid': True,
         'efc__D': True,
         'efc__J': True,
         'efc__J_colind': True,
@@ -1059,12 +1187,25 @@ _BATCH_DIM = {
         'efc__aref': True,
         'efc__force': True,
         'efc__frictionloss': True,
+        'efc__iD': True,
+        'efc__iJ': True,
+        'efc__iJ_colind': True,
+        'efc__iJ_rowadr': True,
+        'efc__iJ_rownnz': True,
+        'efc__iaref': True,
         'efc__id': True,
+        'efc__iforce': True,
+        'efc__ifrictionloss': True,
+        'efc__iid': True,
+        'efc__island': True,
+        'efc__istate': True,
+        'efc__itype': True,
         'efc__margin': True,
         'efc__pos': True,
         'efc__state': True,
         'efc__type': True,
         'efc__vel': True,
+        'efc_islandid': True,
         'energy': True,
         'eq_active': True,
         'flexedge_J': True,
@@ -1073,8 +1214,23 @@ _BATCH_DIM = {
         'flexvert_xpos': True,
         'geom_xmat': True,
         'geom_xpos': True,
+        'history': True,
+        'iqacc': True,
+        'iqacc_smooth': True,
+        'iqfrc_constraint': True,
+        'iqfrc_smooth': True,
+        'island_dofadr': True,
+        'island_efcadr': True,
+        'island_ne': True,
+        'island_nefc': True,
+        'island_nf': True,
+        'island_nv': True,
         'light_xdir': True,
         'light_xpos': True,
+        'map_dof2idof': True,
+        'map_efc2iefc': True,
+        'map_idof2dof': True,
+        'map_iefc2efc': True,
         'mocap_pos': True,
         'mocap_quat': True,
         'moment_colind': True,
@@ -1087,6 +1243,7 @@ _BATCH_DIM = {
         'ne': True,
         'nefc': True,
         'nf': True,
+        'nidof': False,
         'nisland': True,
         'njmax': False,
         'njmax_nnz': False,
@@ -1095,7 +1252,7 @@ _BATCH_DIM = {
         'nworld': False,
         'qLD': True,
         'qLDiagInv': True,
-        'qM': True,
+        'qLU': True,
         'qacc': True,
         'qacc_smooth': True,
         'qacc_warmstart': True,
@@ -1138,9 +1295,23 @@ _BATCH_DIM = {
         'xquat': True,
     },
     'Model': {
+        'D_colind': False,
+        'D_diag': False,
+        'D_rowadr': False,
+        'D_rownnz': False,
         'M_colind': False,
+        'M_elemid': False,
+        'M_fullm_i': False,
+        'M_fullm_j': False,
+        'M_fullm_upper_elemid': False,
+        'M_fullm_upper_i': False,
+        'M_fullm_upper_j': False,
+        'M_mulm_col': False,
+        'M_mulm_madr': False,
+        'M_mulm_rowadr': False,
         'M_rowadr': False,
         'M_rownnz': False,
+        'M_tiles': False,
         'actuator_acc0': True,
         'actuator_actadr': False,
         'actuator_actearly': False,
@@ -1152,6 +1323,7 @@ _BATCH_DIM = {
         'actuator_cranklength': True,
         'actuator_ctrllimited': False,
         'actuator_ctrlrange': True,
+        'actuator_delay': False,
         'actuator_dynprm': True,
         'actuator_dyntype': False,
         'actuator_forcelimited': False,
@@ -1159,6 +1331,8 @@ _BATCH_DIM = {
         'actuator_gainprm': True,
         'actuator_gaintype': False,
         'actuator_gear': True,
+        'actuator_history': False,
+        'actuator_historyadr': False,
         'actuator_lengthrange': True,
         'actuator_trnid': False,
         'actuator_trntype': False,
@@ -1351,6 +1525,8 @@ _BATCH_DIM = {
         'light_poscom0': True,
         'light_targetbodyid': False,
         'light_type': True,
+        'mapD2M': False,
+        'mapM2D': False,
         'mapM2M': False,
         'mat_rgba': True,
         'mat_texid': True,
@@ -1379,6 +1555,7 @@ _BATCH_DIM = {
         'mesh_vertnum': False,
         'mocap_bodyid': False,
         'nC': False,
+        'nD': False,
         'nJfe': False,
         'nJmom': False,
         'nJten': False,
@@ -1403,6 +1580,7 @@ _BATCH_DIM = {
         'ngravcomp': False,
         'nhfield': False,
         'nhfielddata': False,
+        'nhistory': False,
         'njnt': False,
         'nlight': False,
         'nmat': False,
@@ -1479,15 +1657,11 @@ _BATCH_DIM = {
         'pair_solreffriction': True,
         'plugin': False,
         'plugin_attr': False,
+        'qD_fullm_i': False,
+        'qD_fullm_j': False,
         'qLD_all_updates': False,
         'qLD_level_offsets': False,
         'qLD_updates': False,
-        'qM_fullm_i': False,
-        'qM_fullm_j': False,
-        'qM_mulm_col': False,
-        'qM_mulm_madr': False,
-        'qM_mulm_rowadr': False,
-        'qM_tiles': False,
         'qpos0': True,
         'qpos_spring': True,
         'rangefinder_sensor_adr': False,
@@ -1498,9 +1672,13 @@ _BATCH_DIM = {
         'sensor_contact_adr': False,
         'sensor_cutoff': False,
         'sensor_datatype': False,
+        'sensor_delay': False,
         'sensor_dim': False,
         'sensor_e_kinetic': False,
         'sensor_e_potential': False,
+        'sensor_history': False,
+        'sensor_historyadr': False,
+        'sensor_interval': False,
         'sensor_intprm': False,
         'sensor_limitfrc_adr': False,
         'sensor_limitpos_adr': False,