One Step of Beam Search

Implement a single expansion step of beam search for autoregressive decoding.

Signature: def beam_search_step(beams: list, log_probs: np.ndarray, beam_size: int) -> list

• beams: a list of dicts {'tokens': list[int], 'score': float}
• log_probs: shape (len(beams), vocab_size) — next-token log probabilities for each beam
• beam_size: how many beams to keep

For each existing beam i and each vocab token v:

• new_score = beams[i]['score'] + log_probs[i, v]
• new_tokens = beams[i]['tokens'] + [v]

Return the top beam_size candidates ranked by new_score (descending). Break ties by lower beam index, then lower token id.