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, 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?