15. Efficient Batched Beam Search Step

Problem

Beam search keeps the top-k highest cumulative-log-probability sequences at every decoding step instead of just one (greedy) or a full search tree (exhaustive). The naive way to extend a beam is to enumerate every (existing beam, next token) pair, sort all beam_width * vocab_size candidates, and keep the top beam_width — but that full sort is wasteful and doesn't batch cleanly across an inference batch. Production implementations instead flatten the (beam, vocab) candidate grid into one axis and run a single torch.topk over it, which is the standard way to isolate this step as a reusable, batchable primitive.

Given beam_log_probs (shape (batch, beam_width)) and next_token_log_probs (shape (batch, beam_width, vocab)), compute every extended candidate's cumulative log-probability, flatten, and select the beam_width best.

Source: src/15_batched_beam_search.py

def beam_search_step(beam_log_probs, next_token_log_probs, beam_width) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]

>>> import torch
>>> beam_log_probs = torch.tensor([[0.0, float("-inf")]])
>>> next_token_log_probs = torch.log_softmax(torch.rand(1, 2, 4), dim=-1)
>>> vals, beam_idx, token_idx = beam_search_step(beam_log_probs, next_token_log_probs, beam_width=2)
>>> beam_idx  # both new beams must come from beam 0, since beam 1 is -inf
tensor([[0, 0]])

Step-by-Step Approach

1. Broadcast each beam's running cumulative log-probability across all its vocabulary continuations: beam_log_probs.unsqueeze(-1) + next_token_log_probs gives a (batch, beam, vocab) tensor of every candidate's total log-probability.
2. Flatten the (beam, vocab) axes into one with .view(batch, beam * vocab) — this is what turns "sort a beam-by-vocab grid" into "sort one flat list," letting a single torch.topk replace the naive nested sort.
3. Call torch.topk(flat, beam_width, dim=-1) to get the top beam_width candidate log-probabilities and their flat indices in one batched, vectorized call.
4. Undo the flattening with integer division and modulo: beam_idx = top_flat_idx // vocab recovers which of the original beams each winner extends, and token_idx = top_flat_idx % vocab recovers which token was appended.
5. Return the new cumulative log-probabilities alongside beam_idx (needed to reorder any per-beam state like a KV-cache to match the surviving beams) and token_idx.
6. Verify the "nonexistent beam" edge case: seed a beam's log-prob at -inf (as at the very first decoding step, where only one real beam exists) and confirm no winning candidate ever traces back to it — this should fall out of the arithmetic automatically, with no explicit masking code required.

The key insight is that -inf in beam_log_probs is a free, self-enforcing mask: adding -inf to any finite next-token log-probability still yields -inf, so a nonexistent beam's entire row of candidates is automatically ineligible for topk without a single line of special-case logic — representing "this beam doesn't exist" as "arbitrarily bad log-probability" turns a would-be edge case into the natural behavior of the same arithmetic used for every other step.

Reference solution

import torch


def beam_search_step(beam_log_probs, next_token_log_probs, beam_width):
    batch, beam, vocab = next_token_log_probs.shape
    # broadcast each beam's running log-prob across its vocab continuations;
    # a -inf beam poisons its whole row so it can never win a topk below
    candidate_log_probs = beam_log_probs.unsqueeze(-1) + next_token_log_probs  # (batch, beam, vocab)
    # flatten (beam, vocab) into one axis so a single topk replaces a full
    # sort of beam*vocab candidates -- the key efficiency trick
    flat = candidate_log_probs.view(batch, beam * vocab)  # (batch, beam*vocab)
    top_vals, top_flat_idx = torch.topk(flat, beam_width, dim=-1)  # both (batch, beam_width)
    # undo the flattening: integer div/mod recovers which beam and which
    # vocab token each winning flat index came from
    beam_idx = top_flat_idx // vocab
    token_idx = top_flat_idx % vocab
    return top_vals, beam_idx, token_idx

Key Functions & Tricks

• Tensor.unsqueeze(-1) broadcast-add — extends each beam's scalar cumulative log-prob across the whole vocabulary in one op instead of a Python loop over beams.
• Tensor.view(batch, beam * vocab) — the flattening step that turns a 2D-per-batch selection problem into a 1D-per-batch topk call.
• torch.topk(flat, beam_width, dim=-1) — the O(n log k) selection primitive that replaces a full O(n log n) sort of every candidate.
• // vocab and % vocab — integer div/mod recovers the 2D (beam, token) coordinates from a flat index, the inverse of the earlier .view().
• float("-inf") as a "this beam doesn't exist" sentinel — avoids explicit masking logic by making the arithmetic self-enforcing.

How to Recognize This Pattern