Fix prompt lookup decoding length cap
Author: Sunt-ing
Pull request: yes

<!-- ci-dashboard-badge:start -->
[![CI](https://transformers-ci.lor-e.huggingface.cool/badge/pr?pr=47195)](https://transformers-ci.lor-e.huggingface.cool/d/pytest-observability-by-pr/pytest-observability-branch?var-pr=47195)
<!-- ci-dashboard-badge:end -->

# What does this PR do?

`generate(..., prompt_lookup_num_tokens=3, max_new_tokens=2)` can return more than two new tokens. With `distilgpt2`, current main returns length 20 for a target length of 19. Greedy decoding returns length 19, and the overlapping prompt-lookup prefix is identical to greedy.

The length cap in `PromptLookupCandidateGenerator` compares slice indices into the previous `input_ids` with the generated sequence's `max_length`. Since `input_length <= max_length` during generation, that cap does not bind. This PR caps prompt-lookup candidates by the remaining generation budget, with the same one-token reserve used by assisted generation.

This supersedes the closed PR #46993 with a public-checkpoint end-to-end repro. 

<details>
<summary>End to end repro with distilgpt2 and Checks</summary>

Environment:

```text
Hardware: AutoDL c4090, CPU-only run with CUDA_VISIBLE_DEVICES=""
Python: 3.12.3
PyTorch: 2.8.0+cu128
Transformers main: 187fda1bb179f5a2e261736bb4d91c292864ca0d
Model: distilgpt2
Command: CUDA_VISIBLE_DEVICES="" PYTHONPATH=src:. python repro_distilgpt2_prompt_lookup.py
```

Script:

```python
import os

os.environ["CUDA_VISIBLE_DEVICES"] = ""

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer


MODEL_ID = "distilgpt2"
PROMPT = "The quick brown fox jumps over the lazy dog. The quick brown fox jumps over the"
MAX_NEW_TOKENS = 2
PROMPT_LOOKUP_NUM_TOKENS = 3


tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=torch.float32).eval()
encoded = tokenizer(PROMPT, return_tensors="pt")

kwargs = dict(
    attention_mask=encoded.attention_mask,
    max_new_tokens=MAX_NEW_TOKENS,
    min_new_tokens=MAX_NEW_TOKENS,
    do_sample=False,
    num_beams=1,
    pad_token_id=tokenizer.eos_token_id,
)

with torch.no_grad():
    torch.manual_seed(0)
    greedy = model.generate(encoded.input_ids, **kwargs)
    torch.manual_seed(0)
    prompt_lookup = model.generate(
        encoded.input_ids,
        prompt_lookup_num_tokens=PROMPT_LOOKUP_NUM_TOKENS,
        **kwargs,
    )

target_len = encoded.input_ids.shape[-1] + MAX_NEW_TOKENS
overshoot = prompt_lookup.shape[-1] - target_len
overlap = min(greedy.shape[-1], prompt_lookup.shape[-1])
overlap_equal = greedy[0, :overlap].tolist() == prompt_lookup[0, :overlap].tolist()

print("TARGET_LEN", target_len)
print("GREEDY_LEN", greedy.shape[-1], "PROMPT_LOOKUP_LEN", prompt_lookup.shape[-1])
print("OVERSHOOT", overshoot, "OVERLAP_EQUAL", overlap_equal)
print("GREEDY_TEXT", repr(tokenizer.decode(greedy[0])))
print("PROMPT_LOOKUP_TEXT", repr(tokenizer.decode(prompt_lookup[0])))

if overshoot > 0:
    raise SystemExit("prompt lookup generated past max_new_tokens")
if prompt_lookup.shape[-1] != target_len:
    raise SystemExit("prompt lookup did not stop at the target length")
if not overlap_equal:
    raise SystemExit("prompt lookup changed the greedy prefix")
```

Current main:

```text
TARGET_LEN 19
GREEDY_LEN 19 PROMPT_LOOKUP_LEN 20
OVERSHOOT 1 OVERLAP_EQUAL True
GREEDY_TEXT 'The quick brown fox jumps over the lazy dog. The quick brown fox jumps over the lazy dog'
PROMPT_LOOKUP_TEXT 'The quick brown fox jumps over the lazy dog. The quick brown fox jumps over the lazy dog.'
prompt lookup generated past max_new_tokens
```

After this PR:

```text
TARGET_LEN 19
GREEDY_LEN 19 PROMPT_LOOKUP_LEN 19
OVERSHOOT 0 OVERLAP_EQUAL True
GREEDY_TEXT 'The quick brown fox jumps over the lazy dog. The quick brown fox jumps over the lazy dog'
PROMPT_LOOKUP_TEXT 'The quick brown fox jumps over the lazy dog. The quick brown fox jumps over the lazy dog'
```

Reverting only this patch makes the same script fail again with `PROMPT_LOOKUP_LEN 20`, `OVERSHOOT 1`.

<summary>Checks</summary>

```text
git diff --check
passed
```

</details>

## Code Agent Policy

- [x] I confirm that this is not a pure code agent PR.

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://huggingface.co/docs/transformers/contributing) and the
      [Pull Request](https://huggingface.co/docs/transformers/pr_checks) checks?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes according to the [guidelines](https://github.com/huggingface/transformers/tree/main/docs)?
- [ ] Did you write any new necessary [tests](https://huggingface.co/docs/transformers/testing)?

## Who can review?

cc @Cyrilvallez

Comments:
<!-- ci-dashboard-recap:start -->

---

### CI recap

**Dashboard:** [View test results in Grafana](https://transformers-ci.lor-e.huggingface.cool/d/pytest-observability-by-pr/pytest-observability-branch?var-pr=47195)
**Latest run:** [29000852914:2](https://github.com/huggingface/transformers/actions/runs/29000852914)
**Result:** `success` | **Jobs:** 15 | **Tests:** 170,702 | **Failures:** 0 | **Duration:** 19h 41m

<!-- ci-dashboard-recap:end -->