perf(rbd): flat 1-D narrow-phase dispatch — pack warps across batches

[haixuanTao]

Problem

The narrow-phase kernels dispatch [max_len/64, num_batches, 1]: one 64-lane workgroup per batch, rounded up from that batch's live pair count. A batched robot environment has ~7 collision pairs, so at 2048+ envs the GPU runs thousands of workgroups at ~11% lane occupancy, and narrow-phase time scales almost linearly with env count (0.02 ms at 1 env → 18.6 ms at 4096 on the scene below) — it goes from negligible to the single largest pass in the frame.

Change

A new one-thread kernel (gpu_flatten_batches_dispatch, same style and cost as the existing max-scan init kernels) builds exclusive prefix offsets over the per-batch work-lists plus a flat [total/64, 1, 1] indirect grid. The classify, deferred, and PFM kernels walk 0..total and recover (batch, item) with a binary search over the offsets (≤12 steps at 4096 batches). Warps fill with real pairs from consecutive batches instead of one batch's handful plus idle lanes.

Buffer layout is unchanged — per-batch capacity-strided storage stays as is; only the dispatch shape and index math moved. The same flatten kernel is reused for both the collision-pair and the PFM work-lists, repurposing the two existing indirect buffers (nothing else consumes their batched form).

Results

RTX 5090, batched 12-DOF biped (one robot per env, mesh feet), dt = 5 ms, steady-state contacts, wgpu/Vulkan:

	2048 envs	4096 envs
narrow-phase	10.31 → 5.09 ms (2.0×)	18.57 → 6.94 ms (2.7×)
whole step (wall)	22.45 → 16.50 ms	—
throughput	91k → 124k env-steps/s

Physics is bit-identical before/after (robot trajectory and cube-settle checks match to printed precision), as expected — the same items are processed, only warp packing changes. Single-env scenes see no change (the flat grid degenerates to the old one).

The remaining narrow-phase cost is the per-pair GJK/EPA divergence itself; the same per-batch dispatch pattern also exists in the contact-solver kernels (contacts_grid), which this PR deliberately leaves untouched — planned as a follow-up.

🤖 Generated with Claude Code

[haixuanTao]

Pushed an additional commit with a correctness fix found while stress-testing heterogeneous environments: the classify and pfm_pfm kernels read collider_parent without the per-batch stride, resolving every batch's parents through batch 0's table. Invisible with identical envs (the tables coincide — which is why it survived all the bit-exactness checks), silently wrong with per-env MJCF heterogeneity: same-body pairs get skipped/mis-solved against the wrong topology. The buffer's own construction comment ('Env-local body slot; the kernels apply the per-batch stride') documents the intended contract. The bug predates this PR — the unsliced reads exist on main — but since this PR rewrites those exact lines, carrying the fix here avoids a conflicting follow-up. Repro: two envs with equal body/collider counts but different parent tables (glued-boxes vs separate-boxes); pre-fix the dropped box never stacks, post-fix it settles at the expected height. Same commit also cherry-picked onto #22.