I partially solved the issues this way:
for the first problem, I could still not find a truly satisfying solution, but I managed to remove the second for loop and reconvert it into a series of vectorized operations.

avoid computing the gradients for the bio input
compute the mention spans
use torch.cat() to combine the nonzero() result with the right end
compute the gradient on the mentions

The second issue become a bit easier to solve because we already have the mentions in a matrix format. However, some shape calculations were a bit tricker and I had to make use of torch.bmm for the implementation of top-k during inference.
    def forward(
        self,
        X,
        bio_output: Optional[torch.LongTensor],
        entities_output: Optional[torch.LongTensor],
        k=100
    ) -> (torch.tensor, torch.tensor):
        """
        :param x the (raw) output of the first transformer block. It has a shape:
                B x N x (embed_size).
        :param bio_output the output of the bio classifier. If not provided no loss is returned
                (which is required during a training stage).
        :param entities_output the detected entities. If not provided no loss is returned
                (which is required during a training stage).
        
        
        :returns a pair (loss, transformer_output). If either of entities_output or bio_output is
                  None loss will be None as well.
        """


        y = torch.zeros_like(X).to(DEVICE)

        # Disable gradient calculation for BIO outputs, but re-enable them
        # for the span
        with torch.no_grad():

            calculate_loss = bio_output is not None and entities_output is not None
            if calculate_loss:
                loss = torch.zeros((1,)).to(DEVICE)
            else:
                loss = None

            begin_positions = torch.nonzero(bio_output == self.begin)

            # if no mentions are detected skip the entity memory.
            # TODO: would be nice to assess how often this happens.
            if len(begin_positions) == 0:
                return loss, y

            # FIXME: Not really parallelized (we don't have vmap yet...)
            end_positions = torch.tensor([
                self._get_last_mention(bio_output, pos) for pos in begin_positions]).unsqueeze(1).to(DEVICE)

        # Create an array of:
        # 3 dimensions:
        # [ batch_idx1, batch_idx2, batch_idx3... ]
        # [ start_idx1, start_idx2, start_idx3... ]
        # [ end_idx1, end_idx2, end_idx3 ]

        positions = torch.cat([begin_positions, end_positions], 1).T

        first = X[positions[0], positions[1]]
        second = X[positions[0], positions[2]]

        mention_span = torch.cat([first, second], 1).to(DEVICE)
        pseudo_entity_embedding = self.W_f(mention_span) # num_of_mentions x d_ent

        # During training consider the whole entity dictionary
        # Not sure why Pylint thinks self.train is a constant
        # pylint: disable=using-constant-test
        if self.train and bio_output is not None and entities_output is not None:
            alpha = F.softmax(
                pseudo_entity_embedding.matmul(self.E.weight), dim=1)

            # shape: B x d_ent
            picked_entity = self.E(alpha)
        
        else:
            # K nearest neighbours
            topk = torch.topk(self.E.weight.T.matmul(
                pseudo_entity_embedding.T), k, dim=1)

            alpha = F.softmax(topk.values, dim=1)

            # mat1 has size (M x d_ent x k), mat2 has size (M x k x 1)
            # the result has size (M x 256 x 1). Squeeze that out and we've got our
            # entities of size (M x 256)
            picked_entity = torch.bmm(self.E.weight[:, topk.indices].swapaxes(0, 1),
                                      alpha.view((-1, k, 1))).squeeze()

        y[positions[0], positions[1]] = self.W_b(picked_entity)

        # Compared to the original paper we use NLLoss.
        # Gradient-wise this should not change anything
        if calculate_loss:
            loss = self.loss(alpha, entities_output[positions[0], positions[1]])
        else:
            loss = None
        return loss, y