perf: add CUDA graph capture to mjwarp backend - #1273
Merged
Conversation
This was referenced Aug 24, 2026
TATP-233
added a commit
that referenced
this pull request
Aug 24, 2026
7 tasks
TATP-233
added a commit
that referenced
this pull request
Sep 4, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
MjwarpBackend冷路径检查 CUDA driver >= 12.4 与 Warp mempool,并原子捕获固定地址的step、forward、reset_dataCUDA graphs。step按现有nsteps重放 graph;selected-row reset 仍严格保持reset_data -> qpos/qvel upload -> forward -> sync -> host cache顺序。RuntimeWarning并整体回退到原 eager 路径,不留下部分启用状态。SimBackend、env、runner、learner、配置或其他 backend;本 PR 是 Roadmap: 借鉴 mjlab 的性能机制与开发体验优势(7 个候选 sub-issue) #1252 任务1的 implementation child Perf: mjwarp 后端引入 CUDA graph capture #1271。Contract / architecture
SimBackend的输入、输出、timing keys 与 reset/step lifecycle 不变。assign更新已分配数组,不替换 graph 捕获的地址;代码明确要求未来若 owner layer 替换 model/data array,必须重新调用冷路径 capture。Throughput
同机、同任务、同参数分别测量 base
60c2ce7c与本 PR653cd7a1:warp-lang==1.14.0;mujoco-warp==3.10.0.3;mempool enabledMjwarpBackend, G1scene_flat.xml, 2048 worlds,sim_dt=0.02/3, 3 physics substeps/control step,nconmax=128,njmax=256backend.step()wall time and includes control H2D、physics sync、qpos/qvel/sensor D2H cache refreshThe benchmark intentionally reports the production backend barrier rather than the benchmark-only mjwarp adapter. Each control step advances
2048 * 3physics world-steps.Validation
make test-alluv run pytest -q tests/base/test_mjwarp_cuda_graph.py tests/base/test_mjwarp_identity.py— 12 passeduv run --extra mjwarp pytest -q -m slow tests/base/test_mjwarp_backend.py tests/base/test_mjwarp_capabilities.py tests/base/test_mjwarp_differential.py tests/base/test_backend_conformance.py— 12 passed, 6 deselectedCloses #1271
Related to #1252