`generate()` produces incoherent output when `inputs_embeds` has length 1
Author: tyarkoni
Pull request: no
Labels: name: WIP

name: bug

### 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?: <fill in>
- Using GPU in script?: <fill in>
- GPU type: NVIDIA H100 PCIe

### Who can help?

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

### Information

- [ ] The official example scripts
- [x] My own modified scripts

### Tasks

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

### Reproduction

#### Steps to Reproduce

```python
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):

```python
# 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

Comments:
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):

```python
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.

```python
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  

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

a new generation edge case 😓 cc @gante 

Thanks a lot :)

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

https://github.com/huggingface/transformers/blob/5689dd6b8ed0257c6c43c0937d0d224d05774ecc/src/transformers/generation/utils.py#L623

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 https://github.com/huggingface/transformers/pull/40652 so the fix will be quite straightforward. I will make a PR shortly

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.