Implementation:Google deepmind Mujoco JAX Experimental FFI

Knowledge Sources	Google_deepmind_Mujoco
Domains	Foreign Function Interface, JAX Integration, CUDA, Graph Capture, GPU Computing
Last Updated	2026-02-15 04:00 GMT

Overview

Core FFI module that enables NVIDIA Warp kernel execution within JAX computation graphs, supporting multiple CUDA graph capture modes and thread-safe kernel registration.

Description

ffi.py is the primary implementation of the JAX-Warp Foreign Function Interface, replacing the legacy custom_call approach for JAX >= 0.5.0. It defines GraphMode (an IntEnum with NONE, JAX, WARP, WARP_STAGED, WARP_STAGED_EX) and ModulePreloadMode for controlling how Warp modules are loaded onto GPU devices. The module maintains thread-safe global registries for FFI kernels (_FFI_KERNEL_REGISTRY), differentiable kernels (_FFI_DIFF_KERNEL_REGISTRY), callables (_FFI_CALLABLE_REGISTRY), and callbacks (_FFI_CALLBACK_REGISTRY), all protected by locks for concurrent access from XLA execution threads. It builds upon the XLA FFI ctypes bindings from xla_ffi.py.

Usage

This module is the default JAX-Warp interoperability layer as of Warp 1.10. It is used by MJX's FFI helper (mujoco.mjx.warp.ffi) and directly by any code that needs to launch Warp kernels from within JAX computation graphs. The GraphMode enum controls whether CUDA graphs are captured by JAX, Warp, or with staging buffer variants for optimized memory transfers.

Code Reference

Source Location

Repository: Google_deepmind_Mujoco
File: mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/ffi.py
Lines: 1-1614

Key Functions

class GraphMode(IntEnum):
    NONE = 0            # don't capture a graph
    JAX = 1             # let JAX capture a graph
    WARP = 2            # let Warp capture a graph
    WARP_STAGED = 3     # Warp graph with staging buffers, copy inside graph
    WARP_STAGED_EX = 4  # Warp graph with staging buffers, copy outside graph

class ModulePreloadMode(IntEnum):
    NONE = 0             # don't preload modules
    CURRENT_DEVICE = 1   # preload on currently active device
    ALL_DEVICES = 2      # preload on all supported devices

class FfiArg:
    """Represents a single FFI argument with name, type, and in/out direction."""

def check_jax_version():
    """Verify JAX >= 0.5.0 is installed."""

def jax_kernel(kernel, ...):
    """Create a JAX primitive from a Warp kernel using XLA FFI."""

Import

from mujoco.mjx.third_party.warp._src.jax_experimental import ffi
from mujoco.mjx.third_party.warp._src.jax_experimental.ffi import GraphMode, jax_kernel

I/O Contract

Inputs

Name	Type	Required	Description
kernel	wp.Kernel	Yes	Compiled Warp kernel to wrap as a JAX primitive
graph_mode	GraphMode	No	CUDA graph capture strategy (default: NONE)
module_preload	ModulePreloadMode	No	When to preload Warp modules onto GPU devices

Outputs

Name	Type	Description
jax_primitive	Callable	JAX-callable function that dispatches to the Warp kernel via XLA FFI
FfiKernel	registered object	Kernel entry in the global FFI registry for reuse and graph capture

Related Pages

Google_deepmind_Mujoco_JAX_Experimental_XLA_FFI - Low-level XLA FFI ctypes bindings imported by this module
Google_deepmind_Mujoco_JAX_Experimental_Custom_Call - Legacy custom call implementation replaced by this module
Google_deepmind_Mujoco_MJX_Warp_FFI - Higher-level MJX FFI layer that uses this module
Google_deepmind_Mujoco_MJX_Warp_Types - Imports GraphMode from this module

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment