[Draft/RFC] Native-CUDA (cuda-oxide) backend: device build, fixed-grid dispatch, CUDA-graph capture#13
Draft
haixuanTao wants to merge 4 commits into
Draft
Conversation
- Workspace: patch khal/vortx family to the local path checkouts and disable
khal-std default features (the rust-cuda default pulls cuda_std, which does
not build under cuda-oxide's -Zbuild-std).
- nexus_rbd_shaders3d: add the cuda-oxide feature (khal-derive keys its
entry-point codegen off this exact feature name on the shader crate).
- Shader fixes for cuda-oxide, both patterns ported from the nexus-rl fork:
* generalize the LU / jacobian helpers over MaybeIndexUnchecked so they
accept SmemBuf-wrapped workgroup memory,
* ColBranchless::col_b — panic-free Mat2/Mat3 column access (Mat::col
panics with a formatted message, which GPU kernels cannot carry) — and
switch all 41 call sites.
With this, build_cuda-style scripts produce a sm_120 cubin from nexus-cuda's
shaders (cuda-oxide .ll -> llvm-link libdevice -> llc -> ptxas), consumed via
CUDA_OXIDE_SHADERS_PTX_NEXUS_RBD_SHADERS3D at host build time.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
… grids on CUDA On the CUDA backend every DispatchGrid::Indirect costs a stream synchronize + device->host count read (a full GPU drain) and invalidates CUDA-graph capture. Port the nexus-rl fork's fixed-grid mechanism: a dispatch_grid() helper picks a capacity-based grid when enabled (kernels are dispatched at workgroup granularity and bounds-check against the true count buffer, so over-launch is always correct). 17 sites: coloring (5), narrow-phase (3), solver (8), warmstart (1). Default ON for CUDA, OFF for WebGPU/Metal (native indirect dispatch is free there); NEXUS_FIXED_GRID=1/0 overrides. Cube-drop capture demo, CUDA backend: 5.6 -> 16.4 gen-fps. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
- Both sort paths dispatched count/reduce/scan_add/scatter indirectly via num_wgs/num_reduce_wgs; compute the same grids host-side from the known capacities (exact match to gpu_init_sort_dispatch) as DispatchGrid::Grid. - The per-pass SortUniforms tensors (and n_sort_flat) were re-allocated every sort call; cache them keyed on (max_keys, passes, num_batches, per_batch). cuMemAlloc during CUDA-graph capture invalidates the capture, and the churn was free to remove anyway. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
NexusPipeline::capture_rbd_graph records one frame's rbd_steps_per_frame x step dispatch sequence into a CUDA graph (executing it once); replay_rbd_graph replays it with a single cuGraphLaunch. Python: pipeline.capture_cuda_graph(viewer, state) / pipeline.replay_cuda_graph(). Requires the CUDA backend, stable buffer sizes (capture after warmup; the graph records raw buffer addresses) and fixed-grid dispatch (the CUDA default). Timestamps and auto-resize are skipped while capturing/replaying. Cube-drop capture demo: 16.4 -> 115 gen-fps (vs 92 on WebGPU) — one graph launch replaces the host-side re-encode of ~40 dispatches per frame. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
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.
Draft for discussion — nexus running end-to-end on native CUDA (cuda-oxide) from Python, including the multibody solver. Companions: dimforge/khal#9 (device backend + graph-capture runtime) and dimforge/vortx#9 (vortx-shaders device build).
Results (cube-drop + LeRobot bipedal capture demos, RTX 5080 laptop, 640×480)
Physics verified against WebGPU (same trajectories; first numerical validation of the multibody solver under cuda-oxide).
The four commits
nexus_rbd_shaders3dgains thecuda-oxidefeature; two shader-code fixes (LU/jacobian helpers generalized overMaybeIndexUncheckedforSmemBufworkgroup memory; panic-freeColBranchless::col_breplacingMat::col, whose formatted panic can't exist in a kernel). Produces a sm_120 cubin via cuda-oxide → llvm-link(libdevice) → llc → ptxas, consumed throughCUDA_OXIDE_SHADERS_PTX_NEXUS_RBD_SHADERS3D.DispatchGrid::Indirectcosts a stream synchronize + host count read and forbids graph capture; capacity-based grids replace them (kernels bounds-check the true count, so over-launch is safe). Default ON for CUDA only;NEXUS_FIXED_GRIDoverrides. 5.6 → 16.4 fps alone.SortUniformstensors (they were re-allocated every call —cuMemAllocinvalidates a capture).NexusPipeline::capture_rbd_graph/replay_rbd_graph(+ Pythoncapture_cuda_graph/replay_cuda_graph) — onecuGraphLaunchreplays the wholerbd_steps_per_frame × stepsequence. Capture after warmup (records raw buffer addresses); timestamps/auto-resize skipped while replaying.Known discussion points
[patch]block (khal/vortx checkouts) because the published khal-std 0.2.1 / vortx-shaders 0.3.0 lack thecuda-oxidefeature —cargo metadatafails on a clean checkout until khal#9 / vortx#9 are released. Happy to rework once those land..socomes fromhaixuanTao/cuda-oxidebranchfeat/nexus3d-vortx-native-cuda.auto_resize_buffersif preferred.🤖 Generated with Claude Code