non-deterministic project gaussians output

[nikmo33]

I am finding that I am getting severly different outputs for outputs of _ProjectGaussians function when called with certain input values.
The weirder thing is that they seem to be different by various amounts for different inputs. I have other inputs that cause different percentage mismatches shown in the print. Note: The differences do not seem to be coming from floating point differences, rather the resulting outputs are hugely different. See example with radii below
I have spent time investigating the issue with no luck. I have tried using cuda streams for synchronisation, checking for un-initialised memory, detaching etc,

Thought I would share in case anyone else has any ideas.

Here is a minimal example I created to reproduce the issue.

I am currently on commit fecca4f (which is the latest commit as of this post)

import torch
import tyro
from gsplat.project_gaussians import _ProjectGaussians


def repro_non_determinism(iterations: int = 10, device: str ="cuda:0"):
    camera_params = torch.load("examples/params.pt").detach().to(device)
    fx, fy, cx, cy = camera_params[0].item(), camera_params[1].item(), camera_params[2].item(), camera_params[3].item()
    H, W, block_size = camera_params[4].item(), camera_params[5].item(), camera_params[6].item()
    means = torch.load("examples/means3d.pt").detach().to(device)
    scales = torch.load("examples/scales.pt").detach().to(device)
    quats = torch.load("examples/quats.pt").detach().to(device)
    viewmat = torch.load("examples/viewmat.pt").detach().to(device)
    projmat = torch.load("examples/projmat.pt").detach().to(device)
    outputs = []
    for i in range(iterations):
        xys, depths, radii, conics, compensation, num_tiles_hit, cov3d = _ProjectGaussians.apply(
            means.contiguous(),
            torch.exp(scales).contiguous(),
            1,
            quats / quats.norm(dim=1, keepdim=True),
            viewmat.contiguous(),
            projmat.contiguous(),
            fx,
            fy,
            cx,
            cy,
            int(H),
            int(W),
            int(block_size),
        )
        torch.cuda.synchronize()
        outputs.append((xys, depths, radii, conics, num_tiles_hit, cov3d))
    strings = ["xys", "depths", "radii", "conics", "num_tiles_hit", "cov3d"]
    print("Testing outputs")
    i = 0
    for j in range(0, iterations):
        for k in range(6):
            print(
                f"% of elements different for {strings[k]} for idx={i} and idx={j}",
                f" = {round(100*(outputs[i][k] != outputs[j][k]).sum().item() / outputs[i][k].numel(), 2)}%"
            )

if __name__ == "__main__":
    tyro.cli(repro_non_determinism)

which gives me the following output -

Testing outputs
% of elements different for xys for idx=0 and idx=0  = 0.0%
% of elements different for depths for idx=0 and idx=0  = 0.0%
% of elements different for radii for idx=0 and idx=0  = 0.0%
% of elements different for conics for idx=0 and idx=0  = 0.0%
% of elements different for num_tiles_hit for idx=0 and idx=0  = 0.0%
% of elements different for cov3d for idx=0 and idx=0  = 0.0%
% of elements different for xys for idx=0 and idx=1  = 0.0%
% of elements different for depths for idx=0 and idx=1  = 0.0%
% of elements different for radii for idx=0 and idx=1  = 65.61%
% of elements different for conics for idx=0 and idx=1  = 65.61%
% of elements different for num_tiles_hit for idx=0 and idx=1  = 65.61%
% of elements different for cov3d for idx=0 and idx=1  = 65.61%
% of elements different for xys for idx=0 and idx=2  = 0.0%
% of elements different for depths for idx=0 and idx=2  = 0.0%
% of elements different for radii for idx=0 and idx=2  = 65.61%
% of elements different for conics for idx=0 and idx=2  = 65.61%
% of elements different for num_tiles_hit for idx=0 and idx=2  = 33.41%
% of elements different for cov3d for idx=0 and idx=2  = 65.61%
% of elements different for xys for idx=0 and idx=3  = 0.0%
% of elements different for depths for idx=0 and idx=3  = 0.0%
% of elements different for radii for idx=0 and idx=3  = 65.61%
% of elements different for conics for idx=0 and idx=3  = 65.61%
% of elements different for num_tiles_hit for idx=0 and idx=3  = 33.41%
% of elements different for cov3d for idx=0 and idx=3  = 65.61%
% of elements different for xys for idx=0 and idx=4  = 0.0%
% of elements different for depths for idx=0 and idx=4  = 0.0%
% of elements different for radii for idx=0 and idx=4  = 65.61%
% of elements different for conics for idx=0 and idx=4  = 65.61%
% of elements different for num_tiles_hit for idx=0 and idx=4  = 33.41%
% of elements different for cov3d for idx=0 and idx=4  = 65.61%
% of elements different for xys for idx=0 and idx=5  = 0.0%
% of elements different for depths for idx=0 and idx=5  = 0.0%
% of elements different for radii for idx=0 and idx=5  = 65.61%
% of elements different for conics for idx=0 and idx=5  = 65.61%
% of elements different for num_tiles_hit for idx=0 and idx=5  = 33.41%
% of elements different for cov3d for idx=0 and idx=5  = 65.61%
% of elements different for xys for idx=0 and idx=6  = 0.0%
% of elements different for depths for idx=0 and idx=6  = 0.0%
% of elements different for radii for idx=0 and idx=6  = 65.61%
% of elements different for conics for idx=0 and idx=6  = 65.61%
% of elements different for num_tiles_hit for idx=0 and idx=6  = 33.41%
% of elements different for cov3d for idx=0 and idx=6  = 65.61%
% of elements different for xys for idx=0 and idx=7  = 0.0%
% of elements different for depths for idx=0 and idx=7  = 0.0%
% of elements different for radii for idx=0 and idx=7  = 65.61%
% of elements different for conics for idx=0 and idx=7  = 65.61%
% of elements different for num_tiles_hit for idx=0 and idx=7  = 33.41%
% of elements different for cov3d for idx=0 and idx=7  = 65.61%
% of elements different for xys for idx=0 and idx=8  = 0.0%
% of elements different for depths for idx=0 and idx=8  = 0.0%
% of elements different for radii for idx=0 and idx=8  = 65.61%
% of elements different for conics for idx=0 and idx=8  = 65.61%
% of elements different for num_tiles_hit for idx=0 and idx=8  = 33.41%
% of elements different for cov3d for idx=0 and idx=8  = 65.61%
% of elements different for xys for idx=0 and idx=9  = 0.0%
% of elements different for depths for idx=0 and idx=9  = 0.0%
% of elements different for radii for idx=0 and idx=9  = 65.61%
% of elements different for conics for idx=0 and idx=9  = 65.61%
% of elements different for num_tiles_hit for idx=0 and idx=9  = 33.41%
% of elements different for cov3d for idx=0 and idx=9  = 65.61%

As an example of how different the parameters can be I ran the following snippet -

ipdb> x = outputs[0][2]
ipdb> y = outputs[2][2]
ipdb> torch.argmax((x!=y).float())
tensor(725, device='cuda:0')
ipdb> x[725]
tensor(3, device='cuda:0', dtype=torch.int32)
ipdb> y[725]
tensor(2120, device='cuda:0', dtype=torch.int32)

The files that will be needed to reproduce the issue are attached below.
gsplat_repro_files.zip

Thanks for the help!