transformers-community

/

group-beam-search

like 0

Follow

Transformers Community 115

What this commit changes

This commit implements the shared-prefill optimization discussed in Discussion #5.

Before group beam search starts, generate() expands each input prompt across all beams. Therefore, the first model forward processes num_beams identical copies of every prompt, even though the beams have not diverged yet:

Before:
    first forward:  batch_size × num_beams prompt copies
    later forwards: batch_size × num_beams decoding sequences

This commit removes that redundant work:

After:
    first forward:  batch_size × 1 prompt copies
    later forwards: batch_size × num_beams decoding sequences

In tensor terms, the first model forward changes from an input batch of:

(batch_size × num_beams, prompt_length)

to:

(batch_size, prompt_length)

The beam-search state itself remains fully beam-shaped throughout the search. Only the inputs to the first model forward are reduced. After the first token is selected, the cache is expanded to beam shape through the existing cache-reordering path, and all subsequent decoding proceeds normally.

Implementation details

1. Detect the initial shared-prefill step

The optimization is enabled only on the first decoding iteration:

shared_prefill = (
    is_first_iteration
    and not model.config.is_encoder_decoder
    and not (output_attentions or output_hidden_states or output_logits)
)

The original path is preserved for:

• encoder-decoder models;
• calls returning attentions;
• calls returning hidden states;
• calls returning raw logits.

These cases are deliberately left unchanged because the returned model outputs are expected to preserve the original beam-shaped structure.

output_scores does not disable the optimization because the processed scores are explicitly reconstructed with the normal beam shape.

2. Reduce only the first model forward

The normal input_ids and all beam-search state remain shaped for:

batch_size × num_beams

Only the inputs passed to the model on the first iteration are reduced.

The already beam-expanded input_ids contain consecutive identical copies of each batch item, so:

forward_input_ids = input_ids[::num_beams]

selects one prompt per batch item.

For example:

before: [batch 0 beam 0, ..., batch 0 beam N,
         batch 1 beam 0, ..., batch 1 beam N]

after:  [batch 0 beam 0,
         batch 1 beam 0]

Beam-expanded tensors in model_kwargs are reduced in the same way, but only when their leading dimension exactly matches batch_beam_size:

forward_model_kwargs = {
    key: value[::num_beams]
    if isinstance(value, torch.Tensor)
    and value.ndim > 0
    and value.shape[0] == batch_beam_size
    else value
    for key, value in model_kwargs.items()
}

Non-tensor values and tensors that are not beam-shaped are left unchanged.

The reduced inputs then go through the normal generation path:

model.prepare_inputs_for_generation(...)
model(...)

There is no separate model-specific forward implementation.

3. Reuse the shared first-step logits

Before beams diverge, every beam belonging to the same batch item has the same raw next-token distribution.

With the optimization, the first model output contains:

batch_size × 1

logit rows instead of:

batch_size × num_beams

logit rows.

For each beam group, the common raw distribution is repeated to the shape expected by the existing group-processing code:

outputs.logits[:, -1, :].repeat_interleave(group_size, dim=0)

The ordering matches batch_group_indices: each batch item's shared logits are repeated once for every beam in the current group.

Everything after this remains on the existing path. Each group still independently applies:

• log_softmax;
• logits_processor;
• diversity penalties;
• beam scores;
• top-k candidate selection;
• beam_scorer.process().

Therefore, the optimization shares only the identical raw model forward. It does not share or bypass any group-specific beam-search processing.

4. Keep processed scores beam-shaped

When output_scores=True, the model's first output is smaller, but the returned processed scores still need the normal beam dimension.

The score buffer is therefore explicitly allocated as:

processed_score = outputs.logits.new_zeros(
    (batch_beam_size, outputs.logits.shape[-1])
)

The existing group loop then fills the appropriate beam rows:

processed_score[batch_group_indices] = next_token_scores_processed

This preserves the expected beam-shaped score output while still avoiding the redundant model forward.

5. Expand the shared KV cache through the existing reorder

After the first model forward, the KV cache contains one row per batch item:

batch_size × 1

rather than one row per beam.

The beam search itself remains beam-shaped and computes the normal reordering_indices. On the first optimized iteration, these indices refer to beam rows, while the cache currently contains only batch rows.

The commit maps them back to their shared cache rows:

reordering_indices //= num_beams

For example, all selected beams originating from batch item b map back to shared cache row b.

The existing cache-reordering code then runs unchanged:

model._reorder_cache(...)

or:

model_kwargs["past_key_values"].reorder_cache(...)

Because the reorder indices are beam-shaped, selecting the same shared cache row multiple times naturally expands the cache from:

batch_size × 1

to:

batch_size × num_beams

No separate manual cache-copying path is needed.

After this first reorder, the cache has the normal beam shape. On the next iteration, shared_prefill is false, and all subsequent model forwards use the original batch_size × num_beams decoding path.

Scope of the change

The main invariant of the implementation is:

Keep all beam-search state beam-shaped and optimize only the redundant initial model forward.

The commit does not change:

• beam initialization;
• beam scores;
• beam grouping;
• diversity processing;
• logits processors;
• candidate selection;
• beam scorer behavior;
• stopping criteria;
• finalization;
• subsequent decoding iterations.

The optimization changes only:

1. the batch size of the first decoder-only model forward;
2. the reuse of its identical raw logits across beams;
3. the mapping of first-step reorder indices to the shared cache rows.

Testing

I tested the original implementation against the optimized implementation with:

• model: Qwen/Qwen3-0.6B;
• GPU: NVIDIA L4;
• precision: FP16;
• prompt length: 833 tokens;
• 32 beams;
• 4 beam groups;
• 32 returned sequences;
• 8 generated tokens;
• deterministic generation (do_sample=False);
• cache enabled.

The test performs warmup runs, measures five runs of each implementation, reports the median runtime, and compares the two sets of returned sequences symmetrically.

Test source

import platform
import random
import time
from difflib import SequenceMatcher
from statistics import median

import torch
import transformers
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL = "Qwen/Qwen3-0.6B"
OLD = "transformers-community/group-beam-search"
NEW = "lavrenko/group-beam-search"
BEAMS, GROUPS, REPEATS = 32, 4, 5

nums = random.Random(42).sample(range(100, 1000), 200)
numbers = " ".join(map(str, nums))
prompt = (
    "Choose exactly four numbers you like most from this list:\n\n"
    f"{numbers}\n\n"
    "Output only the four numbers, separated by spaces."
)

tok = AutoTokenizer.from_pretrained(MODEL)
model = AutoModelForCausalLM.from_pretrained(
    MODEL, dtype=torch.float16
).cuda().eval()

messages = [{"role": "user", "content": prompt}]
x = tok.apply_chat_template(messages, add_generation_prompt=True,
    enable_thinking=False, return_dict=True, return_tensors="pt",
).to("cuda")
prompt_tokens = x.input_ids.shape[1]


def run(repo):
    torch.cuda.synchronize()
    start = time.perf_counter()
    with torch.inference_mode():
        y = model.generate(**x, custom_generate=repo, trust_remote_code=True,
            max_new_tokens=8, num_beams=BEAMS, num_beam_groups=GROUPS,
            num_return_sequences=BEAMS, diversity_penalty=1.0,
            do_sample=False, use_cache=True,
        )
    torch.cuda.synchronize()
    elapsed = time.perf_counter() - start
    outputs = tok.batch_decode(y[:, prompt_tokens:], skip_special_tokens=True)
    return outputs, elapsed


def similarity(a, b):
    def score(s, t):
        return SequenceMatcher(None, s, t).ratio()

    return sum(max(score(s, t) for t in b) for s in a) / len(a)


print("ENVIRONMENT")
print("===========")
print(f"python:        {platform.python_version()}")
print(f"transformers:  {transformers.__version__}")
print(f"torch:         {torch.__version__}")
print(f"cuda:          {torch.version.cuda}")
print(f"gpu:           {torch.cuda.get_device_name()}")
print(f"prompt tokens: {prompt_tokens}")

print("\nWARMUP")
print("======")
run(OLD)
run(NEW)
print("done")

old_times, new_times = [], []
for _ in range(REPEATS):
    old, old_time = run(OLD)
    new, new_time = run(NEW)
    old_times.append(old_time)
    new_times.append(new_time)

score = (similarity(old, new) + similarity(new, old)) / 2

old_median, new_median = median(old_times), median(new_times)

print("\nRESULTS")
print("=======")
print(f"old median: {old_median:.3f}s")
print(f"new median: {new_median:.3f}s")
print(f"speedup:    {old_median / new_median:.2f}x")
print(f"similarity: {score:.1%}")

Transformers v4

ENVIRONMENT
===========
python:        3.12.12
transformers:  4.57.6
torch:         2.9.0+cu126
cuda:          12.6
gpu:           NVIDIA L4
prompt tokens: 833

WARMUP
======
The following generation flags are not valid and may be ignored: ['temperature', 'top_p', 'top_k']. Set `TRANSFORMERS_VERBOSITY=info` for more details.
done

RESULTS
=======
old median: 2.487s
new median: 1.132s
speedup:    2.20x
similarity: 100.0%

Transformers v5

ENVIRONMENT
===========
python:        3.12.13
transformers:  5.12.1
torch:         2.11.0+cu128
cuda:          12.8
gpu:           NVIDIA L4
prompt tokens: 833

WARMUP
======
done

RESULTS
=======
old median: 2.488s
new median: 1.112s
speedup:    2.24x
similarity: 100.0%

In these tests, the optimized implementation produced approximately a 2.2× speedup in both Transformers v4 and v5, with 100% output similarity and no observed change in the generated results.

The algorithm is intended to preserve the generation result. However, exact bitwise identity cannot be guaranteed in every environment. Changing the first forward from a larger batch containing repeated inputs to a smaller batch can sometimes introduce very small floating-point differences because of numerical rounding or different low-level kernel behavior. In an extremely close beam decision, such a difference could theoretically affect the selected sequence even though the beam-search algorithm and scoring logic are unchanged.

Oke, this would work for most models but not really suitable for Qwen-format family. Qwen has a peculiar way to expand inputs so simply cropping off will result in unaligned images-text pair. You can check it by inspecting qwen code in transformers or with batched multi-image inputs

The speed-up thing is great and given that you need it for you research, I can suggest to create a personal repo for "faster_group_beam_search", copy current repo and add your optim on top. Just be aware that the outputs might diverge in some multimodal models

Thanks — you're absolutely correct. I created a batched multi-image Qwen2-VL test specifically to reproduce the case you described, and it exposed exactly the problem you warned about.

With two samples (one image in the first, two images in the second) and 4 beams, the original implementation reaches the first forward with aligned state:

text batch:            8
images referenced:    12
image_grid_thw rows:  12

My original optimization instead produced:

text batch:            2
images referenced:     3
image_grid_thw rows:  12

So input_ids[::num_beams] reduced the text side while Qwen's specially expanded visual state remained beam-expanded, creating an unaligned image-text pair exactly as you said.

I think the right conclusion is that shared prefill should only run when the beam-expanded generation state follows contracts that can be safely reversed. Otherwise, it should simply keep the existing beam-shaped path unchanged.

I added two follow-up commits:

• ca939e72 — Guard shared prefill for non-standard input expansion
• 091e2396 — Harden shared prefill eligibility and cache expansion

The fast path now requires:

• the standard Transformers input-expansion behavior;
• the standard generation-state update behavior;
• beam-major tensor inputs;
• evaluation mode;
• and either no cache or a fresh supported DynamicCache.

If any of these conditions is not satisfied, including Qwen-style multimodal expansion, the optimization is skipped and the original implementation is used unchanged.

I also changed cache handling on the successful fast path. Instead of implicitly using cache reordering to expand the shared cache back to beam shape, the code now explicitly calls batch_repeat_interleave(num_beams) and then performs the normal beam reordering.

The reproducer is here:

https://colab.research.google.com/drive/1NbrXcicuJmchgdKuBF44DGLYQTlw6EY8

It shows:

original implementation:  PASS, 8 / 12 / 12
fixed implementation:     PASS, 8 / 12 / 12
old optimized commit:     FAIL, 2 / 3 / 12

So thank you for pointing this out — it was a real flaw in the original patch, and the batched multi-image case was exactly the right test for exposing it.