If I am understanding the problem correctly, this should be implementable using a single matrix multiply.
Assume that y is a tensor of shape (n_head, hidden_size), and that each linear layer takes in hidden_size in-dims and returns output_size out-dims. This means that each y'_i has shape (output_size,).
Instead of using n_head number of linear layers, we can store a single large weight of shape (n_head, output_size, hidden_size) and use do a batched matrix multiply between the weight and the y:
weight = torch.randn(n_head, output_size, hidden_size, requires_grad=True)
y = torch.randn(n_head, hidden_size)

y_prime = torch.matmul(y.unsqueeze(-1), weight).squeeze(-1)

This gets a little weirder but can still be extended to when you have a batch dimension in addition to your head dimension and if you care about the bias in a linear layer.
If you’re adventurous and on a nightly build of PyTorch, torch.vmap 6 can do this without the dimension squeezing business:
from torch._vmap_internals import vmap
weight = torch.randn(n_head, output_size, hidden_size, requires_grad=True)
y = torch.randn(n_head, hidden_size)
y_prime = vmap(torch.matmul)(weight, y)