perf(rbd): one-launch segmented bitonic sort for small per-batch key counts#20
Open
haixuanTao wants to merge 1 commit into
Open
perf(rbd): one-launch segmented bitonic sort for small per-batch key counts#20haixuanTao wants to merge 1 commit into
haixuanTao wants to merge 1 commit into
Conversation
…cities The batched radix path costs ceil(bits/4) passes x 5 kernels plus extra batch-id passes (~55 launches per broad-phase rebuild) even when each batch only holds a dozen colliders. For per-batch capacities <= 128, sort each batch with a single 64-wide workgroup doing a shared-memory bitonic sort instead. Stability: the comparator orders by (key, original_index), so equal keys keep their input order and the output is identical to the stable radix path. NEXUS_SMALL_SORT=0 restores the radix path. Robot-RL biped scene (2048 envs, 14 colliders/batch, RTX 5090): host pipeline time 12.3 -> 9.5 ms/step (12-DOF G1), 21.2 -> 17.1 ms/step (29-DOF G1); training stats bit-identical either path. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016Dbpu5NhqXK2Xt3YzUbJRM
haixuanTao
force-pushed
the
perf/small-batch-sort
branch
from
July 16, 2026 16:04
fcd1ed7 to
7e0eec2
Compare
haixuanTao
added a commit
to haixuanTao/nexus
that referenced
this pull request
Jul 17, 2026
…capture CUDA-graph capture of the physics step was failing with STREAM_CAPTURE_INVALIDATED on the unified stack (production captured fine; these are unified-added regressions from the dimforge#20 small-sort and dimforge#22 fused-color merges). Any allocation (cuMemAlloc) or pageable host→device copy inside the captured region invalidates the capture. Two offenders, localized with a cuStreamGetCaptureInfo probe after each op: 1. Small-sort fast path (radix_sort): allocated a fresh SortUniforms tensor every dispatch (`pass_uniforms.clear(); push(Tensor::scalar(..))`), bypassing the `uniforms_key` cache the radix path already uses for exactly this reason. Now cached/keyed by per_batch_max — allocate once, reuse. 2. num_colors_uniform sync (rbd_step): updated the fused-solver color-count uniform inside step() — first as a realloc, then (mistakenly) as an in-place write_buffer, both illegal during capture. Moved the sync to auto_resize_buffers, which runs OUTSIDE the captured region; max_colors only changes there (the ratchet) and in set_max_colors (setup), so the value is stable across a capture. With these + the 1-D fixed-grid fix (daf5a43), BIPED_GRAPH=1 now captures and replays on the unified stack (verified: 12 iters @2048 g1_29dof_agile terrain full-parity, falls within baseline noise). Unblocks the env-per-lane LᵀDL wall-clock win, whose per-kernel gain was hidden by eager launch overhead. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01U2n9RqmxTJb8UG5d1Sjw4W
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
1 participant
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.
What
The LBVH morton sort currently runs the full radix cascade — ceil(bits/4) passes × 5 kernels each, plus the batch-id flattening passes — even when each batch only holds a handful of keys. For a robot-RL style scene (thousands of environments, a dozen colliders each) that's ~55 kernel launches per step to sort 12-element segments.
This adds a fast path in
RadixSort::dispatch: when the per-batch capacity is ≤SMALL_SORT_MAX(128), a single segmented bitonic-sort launch (one 64-wide workgroup per batch, keys held in workgroup memory) replaces the whole cascade.Correctness
NEXUS_SMALL_SORT=0env-var escape hatch for A/B.Measured
Host encode time for the broad phase 12.3 → 9.5 ms/step on a 2048-env, 12-DOF-robot scene (RTX 5090, WebGPU and CUDA backends alike — the win is launch count, not kernel time).
cargo checkclean onnexus_rbd3dandnexus_rbd2d.🤖 Generated with Claude Code
https://claude.ai/code/session_01U2n9RqmxTJb8UG5d1Sjw4W