Repository: huggingface/transformers
Pull request: Fix bug in prefill_chunk_size that ignores disable_compile flag

This PR fixes a bug in the `prefill_chunking` function where the compilation check is not performed before calling `get_compiled_call()`. 

When using `prefill_chunk_size > 0` in a `GenerationConfig`, the model's forward function is always compiled, even if `disable_compile=True` is specified. This happens because the `prefill_chunking` function directly calls `get_compiled_call()` without checking if compilation should occur:

```python
model_forward = self.get_compiled_call(generation_config.compile_config)
```

Modified to: 

```python
compile_forward = self._valid_auto_compile_criteria(model_kwargs, generation_config)
if compile_forward:
    model_forward = self.get_compiled_call(generation_config.compile_config)
```

This ensures that models aren't compiled when `disable_compile=True is set`.

Also fixed a typo in the error message ("chunkink" -> "chunking"), lol


## Before submitting
- [x] 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, Pull Request section?
_Yes, but this is still my first PR here, hope it's ok_
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)?
- [x] Did you make sure to update the documentation with your changes? 
_As far as I understand it, there's no need to change the documentation_
- [ ] Did you write any new necessary tests? 
_No, tested with a simple generation script using prefill_chunk_size=8 and disable_compile=True. Before the fix, torch.compile was being called despite disable_compile=True. After the fix, no compilation occurs as expected._


## Who can review?

@gante

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker
- vision models: @amyeroberts, @qubvel
- speech models: @eustlb
- graph models: @clefourrier

Library:

- flax: @gante and @Rocketknight1
- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- pipelines: @Rocketknight1
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @zach-huggingface and @SunMarc
- chat templates: @Rocketknight1

Integrations:

- deepspeed: HF Trainer/Accelerate: @SunMarc @zach-huggingface
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc @MekkCyber

Documentation: @stevhliu

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)
- datasets: [different repo](https://github.com/huggingface/datasets)
- diffusers: [different repo](https://github.com/huggingface/diffusers)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Maintained examples (not research project or legacy):

- Flax: @Rocketknight1
- PyTorch: See Models above and tag the person corresponding to the modality of the example.
- TensorFlow: @Rocketknight1

 -->


Changed files:
filename: src/transformers/generation/utils.py
patch: @@ -4907,9 +4907,14 @@ def _prefill_chunking(self, input_ids: torch.LongTensor, generation_config: Gene
         input_chunks = torch.split(input_ids[:, :-1], chunk_size, dim=-1)
 
         if "past_key_values" not in model_kwargs:
-            raise ValueError("Cannot use prefill chunkink without a cache")
+            raise ValueError("Cannot use prefill chunking without a cache")
+
+        model_forward = self.forward
+
+        compile_forward = self._valid_auto_compile_criteria(model_kwargs, generation_config)
+        if compile_forward:
+            model_forward = self.get_compiled_call(generation_config.compile_config)
 
-        model_forward = self.get_compiled_call(generation_config.compile_config)
         attention_mask = model_kwargs.pop("attention_mask", None)
 
         past_length = 0