Yeah so if I chunk it down into slices of 999 with a mask each iteration is less than a millisecond, and after looping over all of them, it results in a factor of 10 reduction in overall time spent (max time 40ms, usually ~10ms).
I found a few magic numbers with SummaryOps.cu:  THRESH_NUMBER_BINS_FOR_MULTI_BLOCK_MEM = 100 and THRESH_NUMBER_BINS_FOR_GLOBAL_MEM = 1000 which I’m not sure why is a thing when the proper checks for whether the bins will fit within shared memory are done anyway, and briefly looking at the kernel itself, I’m not sure why there’s only a limit of 100 for smem, I guess to use up all the smem I’d only be able to run one block at a time on an SM, I’m sure a cuda ninja could explain to me why this is the case.
Either way to help anyone else out, here’s my IP:
    def _gen_confusion_mat(self, prediction: torch.Tensor, target: torch.Tensor,
                           mask: torch.Tensor = None) -> torch.Tensor:
        if mask is not None:
            temp = self._n_classes * target[mask] + prediction[mask]

        else:
            temp = self._n_classes * target + prediction

        i = 0
        conf_mat = torch.zeros(self._n_classes**2, dtype=torch.int32, device=target.device)
        while i < self._n_classes**2:
            t_mask = temp >= i
            t_mask &= temp < i + 999
            minlength = self._n_classes**2-i if i + 999 > self._n_classes**2 else 999
            conf_mat[i:i+999] = torch.bincount(temp[t_mask] - i, minlength=minlength)
            i += 999

        return conf_mat.reshape(self._n_classes, self._n_classes)