Florian Hahlbohm has released Fused Zip-NeRF, an open source reimplementation of Zip-NeRF built from handwritten fused CUDA kernels.
Zip-NeRF, from Jon Barron, Ben Mildenhall, Dor Verbin, Pratul Srinivasan and Peter Hedman, paired the anti-aliased ray model of Mip-NeRF 360 with Instant NGP hash grids in 2023. Barron released the official JAX code in January 2024, and Fused Zip-NeRF follows that code base closely. Hahlbohm states in the README that Fused Zip-NeRF trains about 6 to 7 times faster and renders about 10 times faster than the official implementation at the same quality.
Fused Zip-NeRF replaces the JAX autograd pipeline with analytic gradients and uses neither tiny-cuda-nn nor PyTorch autograd. The multisampled hash grid lookups run in kernels compiled at runtime through NVRTC, with the small proposal MLPs fused into them in fp32. The main NeRF MLP is a chain of fp16 cuBLASLt matrix multiplies, the losses are computed together with their gradients in per ray kernels, and a single Adam launch updates every parameter buffer. At inference, rays are processed in Morton order so neighboring rays mostly hit the same cache lines.
Validation so far covers one scene, Garden from the Mip-NeRF 360 dataset. On Garden, Fused Zip-NeRF scores 28.22 PSNR and 0.860 SSIM, against 28.22 and 0.863 for the official code and 28.20 and 0.860 reported in the paper. Training took 35 minutes 47 seconds and about 10 GB of VRAM on an RTX 5090, and a 1297x840 test view renders in about 300 ms on the same card. The README's figures for the official code, about 5 hours of training and 0.1 to 0.2 FPS, were measured on an RTX 4090.
Fused Zip-NeRF ships as a method extension to the NeRFICG framework and its CUDA backend, which compiles kernels at runtime from the repository's sources, cannot be used as a standalone PyTorch extension. Its default config follows the paper's protocol of 65,536 rays per batch over 25,000 iterations, which the paper ran on 8 GPUs and Fused Zip-NeRF runs on one, where the official release's single GPU config uses 8,192 rays over 200,000 iterations.
Fused Zip-NeRF is available now on GitHub under the MIT license.
