`generate()` produces incoherent output when `inputs_embeds` has length 1

[tyarkoni]

System Info

• transformers version: 4.57.1
• Platform: Linux-6.8.0-51-generic-x86_64-with-glibc2.35
• Python version: 3.10.18
• Huggingface_hub version: 0.35.3
• Safetensors version: 0.6.2
• Accelerate version: 1.10.1
• Accelerate config: not found
• DeepSpeed version: not installed
• PyTorch version (accelerator?): 2.8.0+cu128 (CUDA)
• Tensorflow version (GPU?): not installed (NA)
• Flax version (CPU?/GPU?/TPU?): not installed (NA)
• Jax version: not installed
• JaxLib version: not installed
• Using distributed or parallel set-up in script?:
• Using GPU in script?:
• GPU type: NVIDIA H100 PCIe

Who can help?

@xadupre (original author of code in question), @zucchini-nlp

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

Steps to Reproduce

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# Load any causal LM
model = AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")
model.eval()

# Create a single embedding (simulating a prefix like a style token)
single_embedding = torch.randn(1, 1, 768)  # [batch=1, seq_len=1, hidden_dim=768]

# Generate with length-1 inputs_embeds
with torch.no_grad():
    outputs = model.generate(
        inputs_embeds=single_embedding,
        max_length=20,
        do_sample=True,
        temperature=1.0,
        pad_token_id=tokenizer.eos_token_id,
    )

# Decode and observe gibberish
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
print("Generated text:", generated_text)
# Output will be incoherent repetitive tokens like "the the the I if the..."

Comparison with working case (length ≥ 2):

# Add a second embedding (e.g., BOS token embedding)
bos_embedding = model.get_input_embeddings()(torch.tensor([[tokenizer.bos_token_id]]))
two_embeddings = torch.cat([single_embedding, bos_embedding], dim=1)  # [1, 2, 768]

# Generate with length-2 inputs_embeds
with torch.no_grad():
    outputs = model.generate(
        inputs_embeds=two_embeddings,
        max_length=20,
        do_sample=True,
        temperature=1.0,
        pad_token_id=tokenizer.eos_token_id,
    )

generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
print("Generated text:", generated_text)
# Output will be coherent

Expected behavior

Expected Behavior

generate() should produce coherent text regardless of whether inputs_embeds has length 1 or length > 1, as long as the embeddings are valid.

Actual Behavior

With length-1 inputs_embeds, generate() produces incoherent, repetitive gibberish that appears to be high-frequency tokens without proper conditioning on previous context.

Suggested Fix

The _cache_dependant_input_preparation method needs to properly handle the transition from inputs_embeds mode to input_ids mode after the first generation step. Specifically:

1. After the first token is generated from inputs_embeds, set inputs_embeds = None for subsequent iterations
2. Or, maintain proper bookkeeping so that the embeddings prefix is correctly tracked throughout the autoregressive loop
3. Or, ensure Exception 4 logic properly handles the case where we've transitioned from embeddings to token IDs

[tyarkoni]

Further information (note: generated by Claude Sonnet 4.5, but I manually verified this):

Bug Location

The issue is in the _cache_dependant_input_preparation method in transformers/generation/utils.py, specifically around lines 472-481.

Root Cause

The method has special handling for inputs_embeds when input_ids.shape[1] == 0 (Exception 4):

if inputs_embeds is not None and input_ids.shape[1] == 0:  # Exception 4
    inputs_embeds = inputs_embeds[:, -cache_position.shape[0] :]
elif (
    inputs_embeds is not None  # Exception 1
    or (cache_position[-1] >= input_ids.shape[1])  # Exception 3
):
    input_ids = input_ids[:, -cache_position.shape[0] :]

The problem:

1. On the first iteration with inputs_embeds of length 1:

   • input_ids.shape[1] == 0 (empty tensor)
   • The code correctly takes Exception 4 path
   • First token is generated successfully

2. On subsequent iterations:

   • input_ids now contains the generated token, so input_ids.shape[1] > 0
   • Exception 4 condition is FALSE
   • The code falls through to Exception 1: inputs_embeds is not None
   • But inputs_embeds is still the original embedding and doesn't correspond to the newly generated input_ids
   • The autoregressive loop loses track of the generated context

When inputs_embeds has length ≥ 2, the logic works because there's proper coordination between the embedding prefix and subsequent token generation.

[i3hz]

def _cache_dependant_input_preparation(
        self,
        input_ids: torch.LongTensor,
        inputs_embeds: Optional[torch.FloatTensor],
        cache_position: Optional[torch.LongTensor],
    ) -> tuple[torch.FloatTensor, torch.LongTensor]:
        """
        Generic cache-dependent input preparation
        The code is put in a separate function to allow granular unit testing
        as it needs a different implementation to be exportable.

        If we have cache: let's slice `input_ids` through `cache_position`, to keep only the unprocessed tokens
        - Exception 1: when passing input_embeds, input_ids may be missing entries
        - Exception 2: some generation methods do special slicing of input_ids, so we don't need to do it here
        - Exception 3: with synced GPUs cache_position may go out of bounds, but we only want dummy token in that case.
        - Exception 4: If input_embeds are passed then slice it through `cache_position`, to keep only the unprocessed tokens and
          generate the first token for each sequence. Later use the generated Input ids for continuation.

        The current implementation does not rely on ``self`` and could be
        a class method. It is left as a standard method to be easily rewritten.
        """
        if inputs_embeds is not None and input_ids.shape[1] == 0:  # Exception 4
            # Initial embeddings case first forward pass with only embeddings
            inputs_embeds = inputs_embeds[:, -cache_position.shape[0] :]
        elif inputs_embeds is not None and cache_position[0] > 0:
            # After first iteration, stop using embeddings and switch to input_ids
            inputs_embeds = None
            input_ids = input_ids[:, -cache_position.shape[0] :]
        elif (
            inputs_embeds is not None  # Exception 1
            or (cache_position[-1] >= input_ids.shape[1])  # Exception 3
        ):
            # Original logic for other cases
            input_ids = input_ids[:, -cache_position.shape[0] :]
        
        return inputs_embeds, input_ids

Is this the correct fix?
When I passed hello to my testing script as input this is the result
I (HelloHelloHelloHelloHelloHelloHelloHelloHelloHelloHelloHelloHelloHelloHello

After modifying the function to this I got
that was an amazing day," he muttered, "We made good calls. Got to do another

If this is the correct fix I would love submit a PR for it. @zucchini-nlp @xadupre

[i3hz]

@vasqu can you please look into this , idk who else to ping . Sorry for the ping

[vasqu]

a new generation edge case 😓 cc @gante

[i3hz]

Thanks a lot :)

[zucchini-nlp]

The root cause is in the way we try to infer prefill vs decoding stages in this line

transformers/src/transformers/generation/utils.py

Line 623 in 5689dd6

if inputs_embeds is not None and len(cache_position) == inputs_embeds.shape[1]:

The logic for inferring prefill vs decoding has always been imperfect and we need to split out prefill to a separate logic. The PR was just merged for it in #40652 so the fix will be quite straightforward. I will make a PR shortly