Skip to content

Use upper Cholesky factors#1355

Open
adenzler-nvidia wants to merge 1 commit into
google-deepmind:mainfrom
adenzler-nvidia:warp-upper-cholesky
Open

Use upper Cholesky factors#1355
adenzler-nvidia wants to merge 1 commit into
google-deepmind:mainfrom
adenzler-nvidia:warp-upper-cholesky

Conversation

@adenzler-nvidia
Copy link
Copy Markdown
Collaborator

@adenzler-nvidia adenzler-nvidia commented May 13, 2026

Summary

  • Use Warp 1.13 upper Cholesky fill mode for tile Cholesky factorization and solves.
  • Store dense Cholesky factors as upper factors and update solver Hessian/qM seeding paths to write upper-triangular storage.
  • Bump the minimum warp-lang dependency to 1.13.

This should be net-same for most cases, and positive for cases where the tile size is a power of 2.

Benchmarks

Clean-machine reverify against stable baseline, three_humanoids with nworld=8192 nconmax=100 njmax=192:

metric baseline current delta
mean run_time, 2 runs 7.549723 7.333066 -2.87%
mean steps/s, 2 runs 1,085,892.4 1,117,605.2 +2.92%
mean solve, 2 runs 657.815 631.126 -4.06%
Nsight total kernel time 1430.310 ms 1281.372 ms -10.41%

Tests

  • uv run ruff check mujoco_warp/_src/block_cholesky.py mujoco_warp/_src/forward.py mujoco_warp/_src/io.py mujoco_warp/_src/smooth.py mujoco_warp/_src/solver.py mujoco_warp/_src/types.py mujoco_warp/_src/smooth_test.py mujoco_warp/_src/solver_test.py mujoco_warp/_src/support_test.py
  • uv run pytest mujoco_warp/_src/solver_test.py::SolverTest::test_qM_fullm_upper_indices_are_row_sorted mujoco_warp/_src/smooth_test.py::SmoothTest::test_factor_m0 mujoco_warp/_src/smooth_test.py::SmoothTest::test_factor_m1 mujoco_warp/_src/smooth_test.py::SmoothTest::test_factor_solve_i0 mujoco_warp/_src/smooth_test.py::SmoothTest::test_factor_solve_i1 mujoco_warp/_src/support_test.py::SupportTest::test_block_cholesky --tb=short -q
  • git diff --check upstream/main

Comment thread mujoco_warp/_src/types.py Outdated
@adenzler-nvidia adenzler-nvidia marked this pull request as draft May 13, 2026 16:11
@adenzler-nvidia adenzler-nvidia marked this pull request as ready for review May 13, 2026 16:17
@thowell
Copy link
Copy Markdown
Collaborator

thowell commented May 13, 2026

since this pr upgrades warp to 1.13.0 we can remove this guard?

if check_version("warp>=1.13.0.dev20260325"):

@adenzler-nvidia
Copy link
Copy Markdown
Collaborator Author

I can remove the guard, are we ok with defaulting to cuBql for all cases?

@thowell
Copy link
Copy Markdown
Collaborator

thowell commented May 13, 2026

@StafaH can you comment on the bvh constructor? thanks!

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