Optimize group beam search by sharing the initial prefill across beams

custom_generate/generate.py

@@ -1,5 +1,6 @@
 import torch
 from transformers import LogitsProcessor, LogitsProcessorList, StoppingCriteriaList, GenerationConfig
+from transformers.cache_utils import DynamicCache
 from transformers.generation.utils import (
     GenerationMixin,
     GenerateBeamDecoderOnlyOutput,
@@ -241,6 +242,34 @@ def _group_beam_search(
     device = input_ids.device
 
     batch_beam_size, cur_len = input_ids.shape
+    cache = model_kwargs.get("past_key_values")
+    use_cache = model_kwargs.get("use_cache", True)
+
+    # Shared prefill is safe only when the generation state follows the
+    # standard contracts that this optimization knows how to reverse.
+    can_share_prefill = (
+        not model.training
+        and type(model)._expand_inputs_for_generation
+        is GenerationMixin._expand_inputs_for_generation
+        and type(model)._update_model_kwargs_for_generation
+        is GenerationMixin._update_model_kwargs_for_generation
+        and (
+            (not use_cache and cache is None)
+            or (
+                use_cache
+                and isinstance(cache, DynamicCache)
+                and not cache.is_compileable
+                and cache.get_seq_length() == 0
+            )
+        )
+        and all(
+            not isinstance(value, torch.Tensor)
+            or value.ndim == 0
+            or value.shape[0] == batch_beam_size
+            for value in model_kwargs.values()
+        )
+    )
+
     # Does not exist anymore in recent versions!
     if hasattr(model, "_get_initial_cache_position"):
         model_kwargs = model._get_initial_cache_position(cur_len, input_ids.device, model_kwargs)
@@ -313,12 +342,37 @@ def _group_beam_search(
             else 1
         )
 
+        # prefill one prompt per batch item before beams diverge
+        # keep the original path when beam expansion cannot be safely undone
+        shared_prefill = (
+            is_first_iteration
+            and can_share_prefill
+            and not model.config.is_encoder_decoder
+            and not (output_attentions or output_hidden_states or output_logits)
+        )
+
+        # keep the beam-shaped search state; reduce only the first model forward
+        forward_input_ids = input_ids
+        forward_model_kwargs = model_kwargs
+
+        if shared_prefill:
+            # undo generate()'s beam expansion for model inputs
+            forward_input_ids = input_ids[::num_beams]
+            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()
+            }
+
         # do one decoder step on all beams of all sentences in batch
         model_inputs = model.prepare_inputs_for_generation(
-            input_ids,
+            forward_input_ids,
             next_sequence_length=next_sequence_length,
             is_first_iteration=is_first_iteration,
-            **model_kwargs,
+            **forward_model_kwargs,
         )
 
         # prepare variable output controls (note: some models won't accept all output controls)
@@ -344,7 +398,10 @@ def _group_beam_search(
             continue
 
         if output_scores:
-            processed_score = torch.zeros_like(outputs.logits[:, -1, :])
+            # scores remain beam-shaped although the first model output is shared
+            processed_score = outputs.logits.new_zeros(
+                (batch_beam_size, outputs.logits.shape[-1])
+            )
         if output_logits:
             # Copy is needed to avoid keeping a hanging ref to outputs.logits which may be very large for first iteration
             # (the clone itself is always small)
@@ -372,9 +429,12 @@ def _group_beam_search(
             # select outputs of beams of current group only
             # No need to clone() the logits here as they will not retain outputs.logits at the end of the loop
             # .float() is needed to retain precision for later logits manipulations
-            next_token_logits = outputs.logits[batch_group_indices, -1, :].to(
-                dtype=torch.float32, device=input_ids.device
-            )
+            # reuse the common raw distribution; diversity is still applied below per group
+            next_token_logits = (
+                outputs.logits[:, -1, :].repeat_interleave(group_size, dim=0)
+                if shared_prefill
+                else outputs.logits[batch_group_indices, -1, :]
+            ).to(dtype=torch.float32, device=input_ids.device)
 
             next_token_scores = nn.functional.log_softmax(
                 next_token_logits, dim=-1
@@ -483,8 +543,16 @@ def _group_beam_search(
         # (that way the memory peak does not include outputs.logits)
         del outputs
 
-        # NOTE: we need to check if `model._reorder_cache` exists for special models like RAG, RecurrentGemma etc.
+        # NOTE: we need to check if `model._reorder_cache` exists for special
+        # models like RAG, RecurrentGemma etc.
         if model_kwargs.get("past_key_values", None) is not None:
+            if shared_prefill:
+                # Expand the shared prefill cache to the normal beam-shaped
+                # state before applying the standard beam reordering.
+                model_kwargs["past_key_values"].batch_repeat_interleave(
+                    num_beams
+                )
+
             if hasattr(model, "_reorder_cache"):
                 model_kwargs["past_key_values"] = model._reorder_cache(
                     model_kwargs["past_key_values"], reordering_indices