Skip to content

Remove CUDA syncs from circular buffer#1031

Merged
kevinzakka merged 2 commits into
mujocolab:mainfrom
rdeits-bd:rd/rm-syncs-from-circular-buffer
May 29, 2026
Merged

Remove CUDA syncs from circular buffer#1031
kevinzakka merged 2 commits into
mujocolab:mainfrom
rdeits-bd:rd/rm-syncs-from-circular-buffer

Conversation

@rdeits-bd

Copy link
Copy Markdown
Collaborator

I was doing some CUDA profiling of our RL environments and found some surprising CUDA syncs associated with actuator delays. I tracked them down to these two places in circular_buffer.py, in which the attempt to skip an unnecessary op actually slowed things down by forcing a CUDA sync. Note that I'm still pretty new to this; I'm just following https://docs.nvidia.com/dl-cuda-graph/torch-cuda-graph/sync-free-code.html

In our internal example, I got a ~3% speedup from this change, though I suspect the effect size will vary a lot based on whether the rest of the MDP is sync-free.

@kevinzakka I'll try to run your benchmarks from #1020

@rdeits-bd

Copy link
Copy Markdown
Collaborator Author

Ok, got some results for the Go1 velocity task. Here's what I get on #1020 for the Go1 task (with no actuator delay, as is the default):

==========================================================================
Summary (all values in env steps per second):
  Physics SPS: sim.step() only (×decimation per env step)
  Env SPS: full env.step() including managers
  Overhead: time spent on non-physics work (observations, rewards, etc.)
==========================================================================
Task                                 Dec  Physics SPS      Env SPS Overhead
--------------------------------------------------------------------------
Velocity-Flat-Go1                      4      641,353      127,292    80.2%

If I set delay_min_lag = delay_max_lag = 1 for the Go1 articulation, then I get:

==========================================================================
Summary (all values in env steps per second):
  Physics SPS: sim.step() only (×decimation per env step)
  Env SPS: full env.step() including managers
  Overhead: time spent on non-physics work (observations, rewards, etc.)
==========================================================================
Task                                 Dec  Physics SPS      Env SPS Overhead
--------------------------------------------------------------------------
Velocity-Flat-Go1                      4      644,432      112,096    82.6%

a pretty significant drop in environment SPS.

With this PR on top of yours, I get:

==========================================================================
Summary (all values in env steps per second):
  Physics SPS: sim.step() only (×decimation per env step)
  Env SPS: full env.step() including managers
  Overhead: time spent on non-physics work (observations, rewards, etc.)
==========================================================================
Task                                 Dec  Physics SPS      Env SPS Overhead
--------------------------------------------------------------------------
Velocity-Flat-Go1                      4      658,344      118,552    82.0%

i.e. recovering about half of the performance loss.

@rdeits-bd rdeits-bd force-pushed the rd/rm-syncs-from-circular-buffer branch from 79a1b4e to 6486621 Compare May 29, 2026 17:34
@rdeits-bd rdeits-bd marked this pull request as ready for review May 29, 2026 17:34
@kevinzakka kevinzakka merged commit 3285a16 into mujocolab:main May 29, 2026
15 of 19 checks passed
Qingxiaoming pushed a commit to Qingxiaoming/mjlab that referenced this pull request Jun 18, 2026
* Remove CUDA syncs from circular buffer

* ruff
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.

2 participants