I eventually managed to solve this problem, @Aktsvigun.
I found it simpler to use length_penalty=0.
I created a tensor which keeps the scores for the beams that are being updated (similar to what you get w/ the input_ids):
running_scores = torch.cat([running_scores[beam_idx, :], beam_scores.unsqueeze(-1)], dim=-1)
Then, when we add the hypothesis of the finished beams in beam_scorer.process, I also send the running_scores. Finally, I concatenate the sumlogprobs (next_score) to the running_scores. Then, when you get the most likely hypothesis in finalize, the tuple contains the running_scores.
To get the logprobs for each token, one would just need to get the consecutive increments (negative here) in running_scores.