Query shared memory per SM from the device for RNEA launch configs

[maxwbuckley]

Summary

DynamicsLaunchCfg assumes a fixed 100 KB of shared memory per SM when estimating how many blocks can co-reside on an SM:

# Total configurable L1/shared-memory per SM (Ampere=128 KB, Hopper=228 KB).
# Used to check how many blocks can co-reside on one SM.
DEFAULT_SM_SHARED_MEM_CAPACITY = 100 * 1024  # conservative for Ampere

That figure is only correct for compute capability 8.6 and 12.x. Per the CUDA Programming Guide, shared memory per SM is 164 KB on 8.0 and 228 KB on 9.0 / 10.x / 11.0, so the occupancy model underestimates resident blocks by up to 2.3x on those parts.

This reads cudaDevAttrMaxSharedMemoryPerMultiprocessor instead, cached per device.

Why it matters

_warp_align_batches scores two candidate batch counts by blocks_per_sm × threads, where blocks_per_sm = sm_capacity // smem_per_block. Because the divisor is fixed, a bigger-than-assumed SM budget changes which candidate wins. It is not a cosmetic constant — the chosen block shape differs for real robot sizes:

num_links	threads/batch	direction	block @ 100 KB	block @ 228 KB
9	1	forward	96	64
9	4	forward	448	416
21	4	forward	192	160
21	4	backward	64	32
31	4	forward	128	96

Safety

The capacity only selects between batch counts that are already valid. calculate_forward_config / calculate_backward_config apply the per-block shared-memory cap (DEFAULT_MAX_SHARED_MEM) and the 1024 thread-per-block limit before calling _warp_align_batches, and that helper only ever returns the warp-aligned candidate or the next one down — both bounded by its input. So no launch configuration can become invalid, and per-block shared memory can only stay the same or shrink. Tests assert this across all three capacities.

When no CUDA device is available the previous 100 KB value is used, so behavior is unchanged on hosts that cannot query the attribute.

Caching

calculate_forward_config runs on every rnea_forward launch, so the attribute lookup is memoized with lru_cache keyed on device ordinal. A test asserts the cache is hit rather than re-entering the driver.

Testing

curobo/tests/_src/curobolib/test_dynamics_launch_cfg.py, 46 tests:

• launch configs stay within the shared-memory and thread limits at 100 / 164 / 228 KB, across 6 robot shapes, forward and backward
• the heuristic never returns more batches than requested
• the queried value agrees with torch.cuda.get_device_properties(...).shared_memory_per_multiprocessor
• the no-CUDA fallback path
• the lookup is cached

Verified on an RTX 5090 (compute capability 12.0), where the device reports exactly 102400 bytes — so this change is a no-op on 8.6/12.x and only alters behavior on 8.0 / 9.0 / 10.x / 11.0.

I do not have a 9.0/10.x/11.0 part to benchmark end-to-end on, so this is submitted as a correctness fix to the occupancy model rather than a measured speedup. Happy to add numbers if someone can run benchmark/inverse_dynamics_kernel_benchmark.py on Hopper/Blackwell.

Deliberately out of scope

Two related things I left alone to keep this reviewable — happy to follow up if you're interested:

1. DEFAULT_MAX_SHARED_MEM = 48 * 1024 (also in kinematics_config.py). Raising this requires cudaFuncSetAttribute(cudaFuncAttributeMaxDynamicSharedMemorySize) on the affected kernels — the pattern already exists in geometry.py::_validate_and_configure_shared_memory — and it trades blocks/SM against batches/block, so it needs measurement rather than a one-line change.
2. Capping blocks_per_sm by cudaDevAttrMaxBlocksPerMultiprocessor (16 on 9.0/10.x/11.0/12.x). The current model can compute an unreachable number of resident blocks for small-shared-memory launches.