diff --git a/mjx/cuda_requirements.txt b/mjx/cuda_requirements.txt index ec304a7f..f51f641b 100644 --- a/mjx/cuda_requirements.txt +++ b/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 diff --git a/mjx/mujoco/mjx/third_party/warp/_src/jax/custom_call.py b/mjx/mujoco/mjx/third_party/warp/_src/jax/custom_call.py index cb2fce91..416905e6 100644 --- a/mjx/mujoco/mjx/third_party/warp/_src/jax/custom_call.py +++ b/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(): diff --git a/mjx/mujoco/mjx/third_party/warp/_src/jax/ffi.py b/mjx/mujoco/mjx/third_party/warp/_src/jax/ffi.py index 8a1d6cad..ef7b0f43 100644 --- a/mjx/mujoco/mjx/third_party/warp/_src/jax/ffi.py +++ b/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, diff --git a/mjx/pyproject.toml b/mjx/pyproject.toml index 89641020..154c3b5b 100644 --- a/mjx/pyproject.toml +++ b/mjx/pyproject.toml @@ -37,7 +37,7 @@ dependencies = [ [project.optional-dependencies] warp = [ - "warp-lang==1.14.0", + "warp-lang==1.15.0", ] dev = [ "isort",