Efficiently computing a batch of results given a batch assignment vector and series of corresponding matrices

Question

I have a 1D tensor of tokens that belong to different batches. The batch sizes here are uneven. Each batch needs to be multiplied with a corresponding weight matrix. My current approach is using a batch pointer vector and a series of distinct weight matrices corresponding to the unique pointers along with a for loop. I want to efficiently compute a result of shape [num_tokens, output_dim], where each weight matrix has shape [input_dim, output_dim]. I also pad the inputs to a multiple of 8 for harnessing NVIDIA Tensor Cores. Here's an example:

import torch

# shape [num_tokens,]
input_dim, output_dim = 4, 8
ptr = torch.tensor([0, 1, 2, 2, 3, 3, -1, -1]) # -1 means padding
features = torch.randn(ptr.shape[0], input_dim)
weights = [torch.randn(input_dim, output_dim) for _ in range(4)]

unique = torch.unique(ptr, sorted=False, return_inverse=False, return_counts=False)
unique = unique[unique != -1] # ignore padding 

results = []

for i in unique:
    split = features[ptr == i, :]
    # pad each split to multiple of 8 for NVIDIA A100
    # repeat pad embedding to desired size
    pad = (
        torch.empty((-split.size(0)) % 8, split.size(-1))
        .uniform_()
        .to(split.device)
    )
    padded_split = torch.cat((split, pad), dim=0)
    attn_mask = torch.cat((torch.ones(split.size(0)), torch.zeros(pad.size(0)))).to(
        torch.bool
    )

    # forward pass
    result = padded_split @ weights[i]
    # strip padding so I can create a 2D result tensor of correct dimension again
    results.append(result[attn_mask, :])

results = torch.cat(results, dim=0)
pad = (
    torch.empty((ptr.size(0) - results.size(0)), results.size(-1))
    .uniform_()
)
# we can now pad results to match the shape of `ptr`
results = torch.cat((results, pad))

The above approach causes a significant slowdown in my code, and occurs in the forward pass of model inference. I suspect that it's because of the padding. I was looking into scatter operations as a solution using ptr as an index vector, but all available methods only support basic reduction methods like sum, mean, max, etc.

What would be a more efficient way of going about this?