My model involves a triangular matrix, where all entries below the diagonal are considered to be zero. If I use a single kernel, I need to store the matrix as a dense, square matrix with zeros, because PyTorch kernels do not operate on “jagged” tensors; if I can use multiple kernels, then I can store the matrix as columns of varying length, without needing to take up memory for the zeros.
I could get into more detail, but that would be way outside the scope of my original question.