transformers-community

/

group-beam-search

like 0

Follow

Transformers Community 115

I might be wrong here, but while benchmarking the current group beam search implementation, I noticed what looks like a potentially large avoidable cost in the first decoding step.

My understanding is that group beam search currently starts from beam-expanded input_ids, so on the first model call the same prompt/prefix may be prefilled separately for every beam. For example, with:

num_return_sequences = 11
beams_per_group = 3
actual beam rows = 33

the initial forward pass seems to process 33 identical copies of the same prefix before cached decoding starts.

After that first step, cached decoding is fine: only the latest token is passed. So this is not the same as the recent Transformers v5 cache correctness issue. My concern is only about the initial full-prefix prefill.

Small benchmark

I ran this in Colab on:

torch=2.11.0+cu128
transformers=5.12.0
GPU=NVIDIA RTX PRO 6000 Blackwell Server Edition
model=Qwen/Qwen3-8B
beams=33

Benchmark cell:

import time, torch
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_ID = "Qwen/Qwen3-8B"
BEAMS = 33
PREFIX_LENS = [512, 1024, 2048, 4096]

device = "cuda" if torch.cuda.is_available() else "cpu"
tok = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID, torch_dtype="auto", trust_remote_code=True
).to(device).eval()

def sync():
    if device == "cuda":
        torch.cuda.synchronize()

def run(ids):
    mask = torch.ones_like(ids)
    sync()
    t0 = time.perf_counter()
    with torch.inference_mode():
        model(input_ids=ids, attention_mask=mask, use_cache=True)
    sync()
    return time.perf_counter() - t0

base = tok.encode("Beam search prefill benchmark. ", add_special_tokens=False)

print("prefix_len,prefill_once_s,prefill_expanded_s,avoidable_s,expanded_over_once")
for n in PREFIX_LENS:
    ids = torch.tensor([(base * (n // len(base) + 1))[:n]], device=device)
    t1 = run(ids)
    tb = run(ids.repeat_interleave(BEAMS, dim=0))
    print(f"{n},{t1:.4f},{tb:.4f},{tb-t1:.4f},{tb/t1:.2f}")

Output:

prefix_len,prefill_once_s,prefill_expanded_s,avoidable_s,expanded_over_once
512,0.2072,0.9402,0.7330,4.54
1024,0.0690,1.9106,1.8416,27.70
2048,0.1234,3.8781,3.7547,31.42
4096,0.2278,8.0236,7.7957,35.22

The exact numbers are noisy, especially for the shorter prefix, but the longer-prefix behavior looks pretty clear: expanded prefill is roughly proportional to the number of beam rows.

Why I think this might matter

For normal generation this may be acceptable because the cost is paid only once. But for workflows that repeatedly run short beam-search probes from long prefixes, this becomes very expensive. In my case, this kind of grouped probe is used to generate a few alternative local continuations, so the expensive part is not the 8 new tokens but repeatedly prefilling the same long prefix across all beam rows.

For example, in one simple four-prefix benchmark using 512/1024/2048/4096 tokens, the expanded prefill alone accounted for about 14 seconds of duplicate work. If the prefix were computed once and the KV cache/logits were expanded to beams afterward, my gut feeling is that the corresponding grouped-beam workload could potentially become several times faster — maybe something like 4–7x in this particular long-prefix scenario, if the implementation overhead is not too high.

Possible optimization

Would it make sense for group beam search to do something like:

1. prefill the unexpanded batch once;
2. expand/reorder the resulting KV cache to batch_size * num_beams;
3. expand the first-step logits/beam state accordingly;
4. continue the existing group beam search loop unchanged?

I may be missing some complication with cache classes, beam scorer state, model-specific prepare_inputs_for_generation, or exact score equivalence. But conceptually, since all beam rows start from the same prefix, it feels wasteful to run the same full-prefix prefill separately for each beam.

The expected requirement would be that this optimization preserves outputs:

old sequences == optimized sequences
old sequence scores approximately equal optimized sequence scores

on deterministic generation.

If this reasoning is correct, this could be a useful follow-up optimization separate from the recent Transformers v5 cached-decoding fix.