Import Warp v1.15.0

PiperOrigin-RevId: 956678131
Change-Id: I09a672b7335a85a17c95773ffc29eb9272744986

mjx/cuda_requirements.txt

@@ -20,8 +20,8 @@ jax-cuda12-pjrt==0.8.3; python_version >= '3.13' \
 jax-cuda12-pjrt==0.5.3; python_version < '3.13' \
     --hash=sha256:04ee111eaf5fc2692978ad4a5c84d5925e42eb05c1701849ba3a53f6515400cc \
     --hash=sha256:c5378306568ba0c81b230a779dd3194c9dd10339ab6360ae80928108d37e7f75
-warp-lang==1.14.0 \
-    --hash=sha256:12656050545cc77bf9b9b155399496c1a6279b5b6c59e407507d6858a2beb4a2 \
-    --hash=sha256:70cd127d0e9109417099649fedf9d00f39f1307ccb7a6e9fb87661337868d7de \
-    --hash=sha256:f482787e8da9c9ef045601fde99095e16d604fbcc3cbb4a1e0cef0769388b316 \
-    --hash=sha256:936b49ec78237f9760e58cbe9c46ee6f4244aefbd62071c4fa9fd3b313dfa878
+warp-lang==1.15.0 \
+    --hash=sha256:06ef42c3ce522749268056653ea253645eb4a96ca164af74ae5113f0d55a6970 \
+    --hash=sha256:3c0711b39ad98ff924d29654b028ec4eaa02330e207ae3efc72445e9cde05b34 \
+    --hash=sha256:95c169f28bd7d6c78ac4ad62e2df1e61a096033748f757157fa4551aed80d010 \
+    --hash=sha256:ee06e2e76bff84088a09e64315e71ed0e55e94c936ec664d7b05ebd56f66a363

mjx/mujoco/mjx/third_party/warp/_src/jax/custom_call.py

@@ -5,7 +5,7 @@ import ctypes
 from functools import reduce
 
 import warp as wp
-from warp._src.context import _build_launch_bounds, type_str
+from warp._src.context import _build_launch_bounds, _raise_cuda_launch_error, _validate_cluster_launch, type_str
 from mujoco.mjx.third_party.warp._src.jax import get_jax_device
 from warp._src.logger import log_warning
 from warp._src.types import array_t, matches_array_class, strides_from_shape
@@ -129,10 +129,25 @@ def _warp_custom_callback(stream, buffers, opaque, opaque_len):
     hooks = kernel.module.get_kernel_hooks(kernel, device)
     assert hooks.forward, "Failed to find kernel entry point"
 
-    # Launch the kernel.
-    wp._src.context.runtime.core.wp_cuda_launch_kernel(
-        device.context, hooks.forward, bounds.size, 0, 256, hooks.forward_smem_bytes, kernel_params, stream, None
-    )
+    # Reject non-cluster-aligned grids with a clear Python error instead of a
+    # cryptic native CUDA error (block_dim=256, max_blocks=0 below).
+    _validate_cluster_launch(hooks.cluster_dim, bounds.size, 256, 0)
+
+    # Launch the kernel (cluster_dim is cached on the hooks at load time).
+    if wp._src.context.runtime.core.wp_cuda_launch_kernel(
+        device.context,
+        hooks.forward,
+        bounds.size,
+        0,
+        256,
+        int(kernel.grid_stride),
+        hooks.cluster_dim,
+        hooks.forward_smem_bytes,
+        kernel_params,
+        stream,
+        None,
+    ):
+        _raise_cuda_launch_error(kernel, device)
 
 
 def _create_jax_warp_primitive():

mjx/mujoco/mjx/third_party/warp/_src/jax/ffi.py

@@ -471,6 +471,8 @@ class FfiKernel:
                         launch_bounds.size,
                         0,
                         256,
+                        int(self.kernel.grid_stride),
+                        hooks.cluster_dim,
                         hooks.forward_smem_bytes,
                         kernel_params,
                         stream,

mjx/pyproject.toml

@@ -37,7 +37,7 @@ dependencies = [
 
 [project.optional-dependencies]
 warp = [
-    "warp-lang==1.14.0",
+    "warp-lang==1.15.0",
 ]
 dev = [
     "isort",