Summary
On RDNA GPUs (wavefront size 32), the backward kernels compute wrong or NaN gradients, so training silently diverges. simple_trainer on Mip-NeRF 360 bicycle plateaus at PSNR ~14 with zero densification (no Gaussians ever split/duplicated), instead of the expected ~25.
The root cause is that several backward kernels hard-code 64 (the MI300X wavefront size). On RDNA the wavefront is 32, so these assumptions break. Forward passes are unaffected, which makes the failure hard to spot: images render fine, but gradients are corrupted.
I have a fix that abstracts the wavefront size and preserves the existing wavefront-64 (MI300X) code path unchanged. With it, bicycle reaches PSNR 24.83 and tests/test_basic.py goes from 81 failures to 105 passed / 9 failed (the 9 are unrelated — nerfacc's CUDA extension can't JIT-build under ROCm).
Environment
- GPU: AMD Radeon 8060S (gfx1151, RDNA 3.5), Ryzen AI Max+ 395 (Strix Halo)
- Wavefront size: 32
- ROCm: 6.4.2 (Fedora
rocm-hip-devel-6.4.2-2.fc43)
- PyTorch: 2.8.0a0 (Fedora
python3-torch)
- OS / kernel: Fedora 43 / 7.1.3-100.fc43.x86_64
- gsplat: release/1.5.3b2
Note: this uses Fedora's stock ROCm packages only — no TheRock nightly, no container.
Reproduction
python examples/simple_trainer.py default \
--data_dir data/360_v2/bicycle --data_factor 4 \
--result_dir results/bicycleBefore the fix: PSNR 14.4, Step N: 0 GSs duplicated, 0 GSs split throughout, and the final GS count ends up below the initial count (only pruning happens, never densification).
Root cause
Three distinct wavefront-64 assumptions, all in backward code:
-
cg::tiled_partition<64> in the rasterize backward kernels (RasterizeToPixels{3DGS,2DGS,FromWorld3DGS}Bwd.cu). tiled_partition<64> fails a static assert on a 32-lane wavefront, so these files don't even compile for RDNA.
-
Manual lane-mask reductions in the projection backward kernels (Projection*.cu). These compute the absolute lane id as threadIdx.x % 64 and iterate lane masks with for (int i = 0; i < 64; ++i). On wavefront 32 the second (and later) wavefronts of a block get lane ids ≥ 32, so my_label_mask & (1ULL << warp_thread_id) is always 0 — those lanes return before contributing, and their gradients are dropped. This silently corrupts v_viewmats, v_means, v_quats, v_scales. It compiles and runs, but the gradients are wrong.
-
tile_size = 8 default in rasterization() (rendering.py). tile_size=8 gives block_size=64, which dispatches the wavefront-64-only rasterize_bs64_* kernel (built around DPP intrinsics and a 1-block = 1-wavefront-64 layout). On RDNA this produces NaN gradients. tile_size=16 avoids the bs64 path and uses the generic kernel.
Fix
- Introduce a
GSPLAT_WAVE64 compile-time macro, set by setup.py from the target arch (gfx9* → 64, else 32), passed to both hipcc and cxx flags.
- Replace hard-coded
64 with warpSize (device) / GSPLAT_WAVE (host) in the lane-mask loops and shared-memory sizing.
- Guard the
bs64 kernel, its DPP helpers, and the tiled_partition<64> path with #if USE_ROCM && GSPLAT_WAVE64, so RDNA falls back to the generic path.
- Default
tile_size to 16.
The wavefront-64 path is unchanged — on MI300X, warpSize == 64 and GSPLAT_WAVE64 == 1, so all existing behavior is preserved.
Results (gfx1151)
|
before |
after |
bicycle PSNR (30k iters, factor 4) |
14.4 |
24.83 |
| final Gaussian count |
~52k (pruned below init) |
7.8M |
| densification |
none |
normal |
tests/test_basic.py |
81 failed |
105 passed, 9 failed (nerfacc env only) |
The fix touches 10 files (+113 −60). I verified it applies cleanly to a fresh release/1.5.3b2 clone and reproduces the results above.
Happy to open a PR if this direction looks acceptable. This should affect all RDNA targets (gfx1030 / gfx1100 / gfx1151 / gfx120x), not just gfx1151.
Summary
On RDNA GPUs (wavefront size 32), the backward kernels compute wrong or NaN gradients, so training silently diverges.
simple_traineron Mip-NeRF 360bicycleplateaus at PSNR ~14 with zero densification (no Gaussians ever split/duplicated), instead of the expected ~25.The root cause is that several backward kernels hard-code
64(the MI300X wavefront size). On RDNA the wavefront is 32, so these assumptions break. Forward passes are unaffected, which makes the failure hard to spot: images render fine, but gradients are corrupted.I have a fix that abstracts the wavefront size and preserves the existing wavefront-64 (MI300X) code path unchanged. With it,
bicyclereaches PSNR 24.83 andtests/test_basic.pygoes from 81 failures to 105 passed / 9 failed (the 9 are unrelated — nerfacc's CUDA extension can't JIT-build under ROCm).Environment
rocm-hip-devel-6.4.2-2.fc43)python3-torch)Note: this uses Fedora's stock ROCm packages only — no TheRock nightly, no container.
Reproduction
python examples/simple_trainer.py default \ --data_dir data/360_v2/bicycle --data_factor 4 \ --result_dir results/bicycleBefore the fix: PSNR 14.4,
Step N: 0 GSs duplicated, 0 GSs splitthroughout, and the final GS count ends up below the initial count (only pruning happens, never densification).Root cause
Three distinct wavefront-64 assumptions, all in backward code:
cg::tiled_partition<64>in the rasterize backward kernels (RasterizeToPixels{3DGS,2DGS,FromWorld3DGS}Bwd.cu).tiled_partition<64>fails a static assert on a 32-lane wavefront, so these files don't even compile for RDNA.Manual lane-mask reductions in the projection backward kernels (
Projection*.cu). These compute the absolute lane id asthreadIdx.x % 64and iterate lane masks withfor (int i = 0; i < 64; ++i). On wavefront 32 the second (and later) wavefronts of a block get lane ids ≥ 32, somy_label_mask & (1ULL << warp_thread_id)is always 0 — those lanes return before contributing, and their gradients are dropped. This silently corruptsv_viewmats,v_means,v_quats,v_scales. It compiles and runs, but the gradients are wrong.tile_size = 8default inrasterization()(rendering.py).tile_size=8givesblock_size=64, which dispatches the wavefront-64-onlyrasterize_bs64_*kernel (built around DPP intrinsics and a 1-block = 1-wavefront-64 layout). On RDNA this produces NaN gradients.tile_size=16avoids the bs64 path and uses the generic kernel.Fix
GSPLAT_WAVE64compile-time macro, set bysetup.pyfrom the target arch (gfx9*→ 64, else 32), passed to both hipcc and cxx flags.64withwarpSize(device) /GSPLAT_WAVE(host) in the lane-mask loops and shared-memory sizing.bs64kernel, its DPP helpers, and thetiled_partition<64>path with#if USE_ROCM && GSPLAT_WAVE64, so RDNA falls back to the generic path.tile_sizeto 16.The wavefront-64 path is unchanged — on MI300X,
warpSize == 64andGSPLAT_WAVE64 == 1, so all existing behavior is preserved.Results (gfx1151)
bicyclePSNR (30k iters, factor 4)tests/test_basic.pyThe fix touches 10 files (+113 −60). I verified it applies cleanly to a fresh
release/1.5.3b2clone and reproduces the results above.Happy to open a PR if this direction looks acceptable. This should affect all RDNA targets (gfx1030 / gfx1100 / gfx1151 / gfx120x), not just gfx1151.