Skip to content

perf: add CUDA graph capture to mjwarp backend - #1273

Merged
TATP-233 merged 1 commit into
mainfrom
perf/issue-1271-mjwarp-cuda-graph
Aug 24, 2026
Merged

perf: add CUDA graph capture to mjwarp backend#1273
TATP-233 merged 1 commit into
mainfrom
perf/issue-1271-mjwarp-cuda-graph

Conversation

@TATP-233

Copy link
Copy Markdown
Collaborator

Summary

  • MjwarpBackend 冷路径检查 CUDA driver >= 12.4 与 Warp mempool,并原子捕获固定地址的 stepforwardreset_data CUDA graphs。
  • 稳态 step 按现有 nsteps 重放 graph;selected-row reset 仍严格保持 reset_data -> qpos/qvel upload -> forward -> sync -> host cache 顺序。
  • 条件不满足或任一 capture 抛错时,带具体原因发出 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

  • 仅改变 mjwarp backend 内部执行机制,SimBackend 的输入、输出、timing keys 与 reset/step lifecycle 不变。
  • 现有 hot path 只通过 Warp assign 更新已分配数组,不替换 graph 捕获的地址;代码明确要求未来若 owner layer 替换 model/data array,必须重新调用冷路径 capture。
  • MuJoCo、Motrix、Drake、macOS 与非 mjwarp entrypoint 行为不变;不新增 support claim、常规 CI 或长期 benchmark infrastructure。
  • 该实现保持 backend capability boundary,不产生新的公共 contract 或 ADR 决策。

Throughput

同机、同任务、同参数分别测量 base 60c2ce7c 与本 PR 653cd7a1

  • GPU: NVIDIA GeForce RTX 4090, 47 GiB, NVIDIA driver 595.84
  • Runtime: CUDA Toolkit 12.9;Warp driver API 13.2;warp-lang==1.14.0mujoco-warp==3.10.0.3;mempool enabled
  • Workload: production MjwarpBackend, G1 scene_flat.xml, 2048 worlds, sim_dt=0.02/3, 3 physics substeps/control step, nconmax=128, njmax=256
  • Measurement: fresh backend per revision, stand state/control, 10 warm-up calls + 40 measured calls;throughput uses mean full backend.step() wall time and includes control H2D、physics sync、qpos/qvel/sensor D2H cache refresh
Metric Before After Difference
Mean full backend step 10.168 ms 4.582 ms -54.9%
Median full backend step 9.706 ms 4.287 ms -55.8%
Mean physics section 9.944 ms 4.392 ms -55.8%
Control-env-steps/s 201,415 446,996 2.22x / +121.9%
Physics-env-steps/s 604,244 1,340,987 2.22x / +121.9%

The benchmark intentionally reports the production backend barrier rather than the benchmark-only mjwarp adapter. Each control step advances 2048 * 3 physics world-steps.

Validation

  • make test-all
    • ruff, mypy, pyright passed
    • 1776 passed, 24 skipped, 273 deselected, 1 xfailed
    • benchmark smoke: module mode 32/33 passed (1 platform-optional skip), script mode 33/34 passed (1 platform-optional skip)
  • uv run pytest -q tests/base/test_mjwarp_cuda_graph.py tests/base/test_mjwarp_identity.py — 12 passed
  • uv 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 deselected

Closes #1271
Related to #1252

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Perf: mjwarp 后端引入 CUDA graph capture

1 participant