rigel

A 3D Gaussian Splatting renderer and trainer with every GPU stage written from scratch in CUDA: EWA projection, tile binning, front-to-back compositing, the hand-derived backward pass for all of it, a fused appearance MLP, and a headless Vulkan interop path that shares the framebuffer between the two APIs. The Python side is a thin torch.autograd.Function over those kernels plus the optimisation loop. Nothing from the reference implementation is vendored; the math is re-derived and checked against independent references and autograd.

gaussians 160000  gpu Tesla T4
  resolution     entries     forward    backward   fwd fps
  720x720        684,271      1.70 ms      7.52 ms     586.6
 1280x720      1,183,134      2.60 ms     15.44 ms     385.3
 1920x1080     2,154,540      4.60 ms     34.12 ms     217.6

^{app/bench.py on the gallery scene. Forward is projection, binning and rasterization; backward is both gradient kernels.}

Gallery

A 160,000-Gaussian procedural scene — a sphere shell coloured by its normals, a three-turn helix, a checkered plane — rendered by the forward pipeline at 360x360, 60 frames, 1.7 ms a frame. It is also the scene the trainer is fitted to below: app/make_dataset.py renders it from 100 random viewpoints and 24 held-out orbit views into the Blender transforms.json layout, so ground truth is exact and every pixel of error is the reconstruction's.

front	orbit	top

The kernels write two per-pixel buffers besides the image, because the backward pass needs them. They are worth looking at.

Image	Accumulated alpha	Overdraw
rasterizer output	1 - final_T, the transmittance left at each pixel	n_contrib, how many Gaussians each pixel examined before terminating

The 16x16 tile grid shows through the overdraw map: binning is per tile, so pixels in one tile see the same candidate list and stop at similar depths. The sphere interior is mottled because the shell is dense and pixels there examine hundreds of candidates before their transmittance drops below threshold; the peak in this view is 1,211. The alpha map is what the backward pass starts from when it walks compositing in reverse.

Reconstruction

Fitting from 50,000 random points, 30,000 iterations, L1 + D-SSIM, adaptive densification until iteration 15,000. Held-out views, 400x400.

scene	appearance	PSNR	SSIM	gaussians	time
orbit	spherical harmonics, degree 3	27.02 dB	0.930	181,555	18.0 min
orbit	fused MLP, 8-d latent	30.18 dB	0.958	251,277	11.0 min
lego (NeRF-synthetic, 400x400)	spherical harmonics, degree 3	33.42 dB	0.983	282,489	26.5 min
lego (NeRF-synthetic, 400x400)	fused MLP, 8-d latent	30.86 dB	0.977	342,601	13.5 min

The two appearance rows differ only in how colour is produced: 48 SH coefficients per Gaussian evaluated in PyTorch, or an 8-float latent decoded by a two-layer MLP in one CUDA kernel. Everything else — the rasterizer, its backward, the densification schedule — is shared, which is what the colour decoupling in the rasterizer was for.

Ground truth	Spherical harmonics	Fused MLP
held-out view 0	27.02 dB · 0.930	30.18 dB · 0.958
held-out view 5	mean over 24 held-out views	mean over 24 held-out views

The MLP row is also the faster one: 21.9 ms an iteration against 36.0 at the end of training. An 8-float latent is a sixth of the optimiser traffic of 48 SH coefficients, and the decoder is one kernel launch over the Gaussians.

The same SH model orbited at radius 4.3 — novel viewpoints, none of them training views. Floaters hang off the checkered plane and the individual Gaussians resolve on the sphere where the shell is thinly sampled. That is what 27 dB looks like on this scene. Held-out PSNR peaked at 27.9 dB at the 10k checkpoint and drifted to 27.0 dB by 30k; the training views are random directions and the test views a fixed orbit, and the model fits the former a little too well once densification stops at 15k.

On real data the picture is cleaner. NeRF-synthetic lego, 100 training views, evaluated on all 200 test views at 400x400 with the alpha composited onto white:

Ground truth	Spherical harmonics	Fused MLP
test view 0	33.42 dB · 0.983	30.86 dB · 0.977
test view 3	282,489 Gaussians, 26.5 min	342,601 Gaussians, 13.5 min

Here PSNR climbs monotonically — 31.8 at 10k, 33.0 at 20k, 33.4 at 30k — with none of the orbit scene's late drift, which points at that drift being a property of the synthetic view split rather than of the optimiser.

Spherical harmonics	Fused MLP
33.42 dB	30.86 dB

Both lego models orbited at radius 4.0 and 0.5 rad elevation — inside the range the training cameras cover, but none of these are training views. The MLP model shows a few soft floaters around the bucket that the SH model does not; that is the 2.6 dB.

The appearance comparison does not have a single winner. On the orbit scene the MLP leads by 3.2 dB; on lego, SH leads by 2.6 dB. The orbit scene's colour is smooth and mostly view-independent, a regime where a shared decoder over an 8-float latent is a strong prior and 48 free coefficients per Gaussian mostly add noise. Lego has sharp, high-frequency albedo across hundreds of thousands of Gaussians, and there the per-Gaussian capacity of SH pays for itself while one small MLP shared by all of them becomes the bottleneck. What is consistent is cost: the MLP runs roughly twice as fast per iteration on both scenes, so it is a real option when a scene tolerates it rather than a strict replacement.

Build

Needs PyTorch with CUDA, a CUDA toolkit whose major version matches it, ninja and pytest. pillow for the gallery tools. Developed against PyTorch 2.8 / CUDA 12.8 / sm_75.

source scripts/env.sh
pip install -e . --no-build-isolation
cd tests && python -m pytest -q

scripts/env.sh pins CUDA_HOME, PATH and TORCH_CUDA_ARCH_LIST. It exists because the development box ships a newer system CUDA ahead of the one PyTorch was built against, and a major-version mismatch fails in unhelpful ways. Edit it for your machine.

Vulkan is optional and detected at build time: if vulkan/vulkan.h is on the include path the interop module is compiled in, otherwise vulkan_available() returns false and everything else works. The loader is dlopened at runtime, so there is nothing to link against.

Running it

Make the synthetic dataset, train on it, evaluate on the held-out views, and spin the result:

python app/make_dataset.py --out data/orbit --res 400 --train 100 --test 24
python app/train.py --data data/orbit --out runs/orbit_sh --iters 30000
python app/train.py --data data/orbit --out runs/orbit_neural --iters 30000 --appearance neural
python app/orbit_gif.py --model runs/orbit_sh/model.pt --out orbit.gif

scene extent 5.370  train views 100  test views 24  gaussians 50000
appearance sh  resolution 400x400
[1000] densify clone 2184 split 2285 prune 1381 -> 40743
[1000] loss 0.1172  gaussians 40743  17.1 ms/it  17s
[5000] densify clone 2036 split 102 prune 106 -> 130459
[10000] loss 0.0156  gaussians 135700  28.3 ms/it  283s
[10000] test psnr 27.92  ssim 0.9515
[20000] test psnr 27.73  ssim 0.9439
[30000] loss 0.0186  gaussians 181555  36.0 ms/it  1080s
final test psnr 27.02  ssim 0.9304  gaussians 181555  time 1082s

Iteration time climbs with the Gaussian count as densification runs, then flattens after it stops at 15,000. The per-iteration loss is a single random training view, so it is noisy by construction; the held-out numbers are the ones to read.

--data takes any directory in the Blender transforms_{train,test}.json layout, so the NeRF-synthetic scenes drop in directly; --white-bg composites their alpha onto white as the reference does, --downscale 2 halves the resolution. train.py writes log.txt, metrics.json, model.pt, a .ply, and prediction/ground-truth pairs under renders/.

Render the gallery and the benchmark table:

python app/render_gallery.py
python app/bench.py

The tree is developed locally and built on a remote GPU. scripts/sync.sh tars the source across and then verifies the remote tree by checksum, because a silently failed sync once made a stale rebuild look like a fix that did not work.

Numbers

Every kernel is compared against an independent reference written before it, and every gradient against PyTorch autograd.

stage	reference	agreement
projection forward	float64 PyTorch	exact cull mask, 1e-5 on values
tile binning	Python re-implementation	exact ordering and per-tile ranges
rasterizer forward	NumPy, same guards	exact contributor counts, 1e-4 image
rasterizer backward	autograd through a differentiable torch loop	2e-7
projection backward	autograd, float64	6e-8 means, 4e-6 scales, 4e-6 quaternions
full pipeline	composition of the two references	2e-6 on all five parameter groups
appearance MLP	float64 torch MLP, three sizes	1e-6 forward, 2e-7 on all five gradients
Vulkan interop	frame read back through a Vulkan transfer queue	bit-identical

Twenty-seven tests, a few seconds on a T4.

$ cd tests && python -m pytest -q -s
  colors    max|ref|=4.1858e+00  rel_err=1.994e-07
  opacity   max|ref|=3.7497e+00  rel_err=1.908e-07
  conic     max|ref|=1.6962e+00  rel_err=4.568e-07
  mean2D    max|ref|=2.3387e+00  rel_err=1.338e-07
  means3D   max|ref|=2.3268e+02  rel_err=6.178e-08
  scales    max|ref|=8.9987e+01  rel_err=4.390e-06
  quats     max|ref|=3.1675e+00  rel_err=3.654e-06

The backward pass is 4.4x to 7.4x slower than the forward and the gap grows with resolution. That is atomicAdd contention: every pixel a Gaussian covers adds into the same handful of gradient slots, and a large on-screen Gaussian serialises thousands of pixels on those addresses. It is the obvious next optimisation and is deliberately left until it can be measured on real training workloads rather than a synthetic scene.

How it works

Five kernels and the CUB primitives between them. Forward runs three; backward runs the other two in reverse.

means3D, scales, quats, camera
        |
  [1] project        -> means2D, depth, radii, conic          per Gaussian
        |
  [2] bin            -> point_list, tile_ranges                per tile, depth sorted
        |
  [3] rasterize      -> image, final_T, n_contrib              per pixel
        |
      loss
        |
  [4] rasterize_bwd  -> dL/d means2D, conic, opacity, colour
        |
  [5] project_bwd    -> dL/d means3D, scales, quats

Projection

Each Gaussian is a mean, an anisotropic covariance parameterised as a rotation and three scales so it stays positive semi-definite, an opacity and a colour. The mean is transformed into camera space and projected through an explicit pinhole (fx, fy, cx, cy). The covariance is projected by linearising the perspective map at the Gaussian's centre:

J = | fx/z    0    -fx x/z^2 |          Sigma_2d = J W Sigma W^T J^T + 0.3 I
    |  0    fy/z   -fy y/z^2 |

x/z and y/z are clamped to 1.3 times the frustum half-width before J is built, because far off-axis Gaussians have a badly wrong linearisation. The 0.3 I is the EWA low-pass filter: it guarantees every Gaussian covers roughly a pixel so nothing falls between samples. The 2x2 is inverted into a conic, and a bounding radius of three standard deviations along the major axis is computed for binning. Anything behind the near plane, degenerate, or entirely off-screen gets radius zero and is dropped by every later stage.

Binning

The image is split into 16x16 tiles and every (Gaussian, tile) pair becomes one 64-bit key:

key = (tile_id << 32) | float_bits(depth)

One cub::DeviceRadixSort over those keys yields tile-major, front-to-back order. The low half works because depth is strictly positive after near-plane culling, and positive IEEE-754 floats have bit patterns monotonic in value, so no float comparator and no second sort are needed. A count kernel, a cub::DeviceScan to give every Gaussian a private slice of the key array (so the expansion needs no atomics), an expansion kernel, the sort restricted to the bits the tile count can occupy, and a range kernel that walks the sorted keys once. The total entry count is the one device-to-host sync in the pipeline.

Rasterization

One block per tile, one thread per pixel. The block cooperatively stages 256 Gaussians at a time into shared memory, then every pixel evaluates them in order:

power = -0.5 (A dx^2 + C dy^2) - B dx dy
alpha = min(0.99, opacity * exp(power))
C += colour * alpha * T;   T *= 1 - alpha

Front-to-back order is what makes early termination legal: once a pixel's transmittance is below 1e-4, nothing behind it can visibly change it, so it stops; once every pixel in the block has stopped, the block exits. __syncthreads_count is already the barrier the loop needs, so the vote is free. Colours arrive precomputed per Gaussian so the kernel is agnostic to how they were produced.

Backward

The forward is C = sum c_i a_i T_i + T_final bg with T_i = prod_{j<i}(1 - a_j), so the gradient with respect to any alpha depends on everything behind it. The backward therefore walks each tile back to front, carrying a suffix accumulation of the colour behind the current Gaussian. Two forward outputs make that possible without storing per-pixel-per-Gaussian state: final_T, undone one step at a time by T_i = T_{i+1} / (1 - a_i), and n_contrib, so entries the pixel never reached are skipped identically. Alpha is clamped at 0.99, which bounds that divisor at 0.01 — the clamp is what keeps the reconstruction numerically viable, not a cosmetic choice.

The projection backward takes dL/d means2D and dL/d conic and chains through the 2x2 inverse, Sigma_2d = T Sigma_3d T^T, T = J W, the Jacobian's dependence on p_cam (with the clamp branch sending gradient to z instead of x when it was active), the pinhole, the camera transform, and finally Sigma_3d = M M^T with M = R(q) S down to scale and quaternion. The quaternion gradient is projected off the radial direction because the forward normalises it.

Deliberate differences from the reference implementation

• Conventions. Row-major world-to-camera matrices and explicit pinhole intrinsics instead of column-major glm and an NDC projection matrix. Datasets hand over intrinsics; the NDC round trip only adds sign errors.
• Alpha clamp gradient. A Gaussian saturated at alpha 0.99 cannot change the image by growing or becoming more opaque, so the gradient through it is zero here. The reference differentiates as though the clamp were absent.
• Quaternion normalisation. Included in the backward, since the forward normalises. The reference assumes unit quaternions. The test feeds non-unit ones to make sure the term is right.
• Colour decoupling. The rasterizer never evaluates spherical harmonics. That is what lets the appearance MLP replace them without touching the compositing kernels.

Appearance MLP

csrc/appearance.cu. Per Gaussian: an 8-float latent concatenated with the unit view direction, a 32-wide ReLU layer, a 3-wide sigmoid layer. Weights live in shared memory; one thread handles one Gaussian. The backward recomputes the hidden layer, accumulates weight gradients into a shared-memory copy with block-local atomics, and flushes each block's partial sums to global memory with one atomic per weight per block — the same reduction pattern the rasterizer backward is missing. Latents get their gradient directly; view directions are constants. Sized for up to 29 latent dims and 64 hidden units.

Vulkan interop

csrc/vk_interop.cpp. libvulkan.so.1 is dlopened and every entry point resolved through vkGetInstanceProcAddr, so nothing is linked. The physical device is matched to the CUDA device by UUID. A VkBuffer is allocated with VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_FD_BIT, the fd exported with vkGetMemoryFdKHR and imported on the CUDA side with cudaImportExternalMemory; the mapped pointer is wrapped as a CUDA tensor and the rasterizer writes into it. To prove the memory is genuinely shared rather than copied, the test renders a frame into it and reads it back through a Vulkan command buffer and a host-visible staging buffer; the two must be bit-identical. Synchronisation is a stream sync before the Vulkan copy, which is correct and coarse; external semaphores are the refinement for a live viewer. Headless only — there is no surface, no swapchain and no window, which is what a T4 in a datacenter can actually run.

Layout

csrc/                  the kernels
  common.cuh             quaternion -> rotation, 3D covariance, EWA projection
  config.h               tile size, low-pass constant, minimum radius
  project.cu             [1]      project_backward.cu   [5]
  binning.cu             [2]      appearance.cu         fused MLP fwd/bwd
  rasterize.cu           [3]      vk_interop.cpp        shared allocation, readback
  rasterize_backward.cu  [4]      ext.cpp               pybind11 module
rigel/                 python
  render.py              Camera, autograd.Function, SH evaluation, render()
  model.py               GaussianModel: activations, densify/prune, save/load
  appearance.py          NeuralAppearance module over the fused kernel
  data.py                Blender transforms.json reader and writer
  loss.py                L1, SSIM, PSNR
app/                   command line
  train.py               fit a scene, evaluate, save
  make_dataset.py        render the procedural scene into a dataset
  render_gallery.py      the still gallery and diagnostics
  orbit_gif.py           spin a model or the procedural scene
  bench.py               the timing table
  run_experiments.sh     every run behind the tables above
tests/                 references and comparisons, one file per stage
scripts/               env.sh toolchain pins, sync.sh verified push
docs/gallery/          the images in this file
data/orbit/            the procedural dataset, 100 train + 24 test views, 400x400
data/nerf_synthetic/   lego from the NeRF-synthetic release, all splits
runs/                  every training run behind the tables above

Data and runs

Everything the tables were computed from is in the tree. data/orbit is the output of app/make_dataset.py and reproduces bit for bit. data/nerf_synthetic/lego is the standard scene, unmodified, via the pablovela5620/nerf-synthetic-mirror dataset on Hugging Face. Each runs/<name>/ holds log.txt, metrics.json, model.pt (loadable with GaussianModel.load), model.ply, and renders/ with prediction and ground-truth pairs at every evaluation.

runs/orbit_sh_flatlr is kept deliberately. It is the SH run from before the higher-order coefficients were given their reduced learning rate: 23.7 dB at 10k falling to 21.3 dB at 30k, against 27.0 dB after the fix. It is the clearest record in the repository of a single constant being worth six decibels.

The model files are 40 to 65 MB each. GitHub accepts them, but they belong in Git LFS if the repository is going to be cloned often.

Tests

cd tests && python -m pytest -q -s

reference.py is a float64, differentiable PyTorch projection written before the kernel. test_project.py checks the kernel against it, including exact agreement on which Gaussians are culled. test_binning.py re-implements the expansion and stable sort in Python and compares the full sorted list and every tile range, then asserts each range is depth-ordered. test_rasterize.py holds a NumPy compositor that reproduces every guard and the contributor counting; counts are compared exactly, because an off-by-one there would not show in the image but would make the backward include or drop a Gaussian. test_rasterize_backward.py and test_project_backward.py run autograd through differentiable references and compare against the hand-derived kernels; the projection test builds a dedicated set of off-axis Gaussians large enough to survive culling so the frustum-clamp branch is actually exercised, because a random scene produced six such cases. test_autograd.py does the same for the composed autograd.Function. test_appearance.py checks the fused MLP forward and all five gradients against a float64 torch MLP at three sizes. test_vulkan.py writes a pattern into the shared allocation from CUDA and reads it back through Vulkan, then does the same with a rendered frame; it skips itself where no ICD is present. test_model.py covers the optimisation bookkeeping: densification and pruning must leave every parameter and its Adam moments the same length, in both appearance modes; the higher-order SH rate is a twentieth of the DC rate; and a model survives a save/load round trip.

Scenes are small on purpose. Gradient correctness needs every code path to fire, not a big scene, and the whole suite runs in a few seconds.