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))
    padded_split = torch.cat((split, pad), dim=0)
    attn_mask = torch.cat((torch.ones(split.size(0)), torch.zeros(pad.size(0)))).to(

    # 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))
# 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?



Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge that you have read and understand our privacy policy and code of conduct.