I tried three more things.
sparse_weights.mm(P)[batch]

as you suggested does indeed work. One epoch now takes ~3min20secs instead of the original 30 seconds, which is still not great, but better than the other options I tried. I have to determine whether this is really worth it.
The approach with
batch_map.smm(sparse_weights).mm(P)

did not work, because apparently .smm / .spmm are no longer available. There exists some external library for handling sparse matrix multiplication, but apparently there is a problem with it properly handling gradients.
One other thing I tried is the following: I realized that of the 30000 groups, only approximately 2000 are occurring in one batch of size ~35000.
So given a batch, I extracted these unique indices into a list, converted this list to a dense one-hot matrix, and multiplied it with the sparse group_weights matrix to get only the relevant weights. The resulting dense matrix multiplied with P gave me the 2000 relevant averaged embeddings, from which I could extract the batch of averaged embeddings corresponding to batch. Basically I though that by not calculating all averaged embeddings I could save time.
Unfortunately, contrary to what I had hoped, this approach took much longer again, with 30minutes per epoch.