tracker: `generate` composability refactor 
Author: gante
Pull request: no
Labels: name: WIP

## `generate` + composability = more use cases with minimal rewrites

As I write this issue, `generate` is mostly a sequential monolith. Many internal blocks were carved into functions over the last two years, but navigating there as a beginner is still messy. It is also very challenging to adapt `generate` to different tasks and/or modalities, forcing us to overwrite the entire generate function (e.g. [RAG](https://github.com/huggingface/transformers/blob/ccdabc5642bf84849af93f591e207dc625c8e1e1/src/transformers/models/rag/modeling_rag.py#L907), [MusicGen](https://github.com/huggingface/transformers/blob/ccdabc5642bf84849af93f591e207dc625c8e1e1/src/transformers/models/musicgen/modeling_musicgen.py#L1542)). All these aspects make using, documenting, maintaining, and testing `generate` a challenge.

This issue is a tracker for the refactor of `generate`, where we aim to build the structure outlined in [this board](https://miro.com/app/board/uXjVNgMGfaQ=/). Key ideas for this refactor:
👉 All models can use the base `generate` API
👉 Reduce if/else blocks
👉 Reduce the barriers to entry for new decoding methods/modalities/use cases
👉 Reduce per-model overwrites when possible
👉 Add unit tests
👉 Add documentation regarding the structure of `generate`

### Tasks

- [ ] 1. Isolate prefill into a separate function, pulling it from the decoding methods. Note that 
    - a) prefill is done excluding the latest token (`input_ids[:, -1:]`), so we don't compute variables regarding the latest token twice; 
    - b) prefill only runs when `use_cache=True` and cache length < input length - 1; 
    - c) `_expand_inputs_for_generation` needs to be changed (it copied inputs before prefill, we will need to copy prefill outputs)
- [ ] 2. (depends on 1.) Separate generate on the 5 stages described in the diagram, passing around the data structures described therein
- [ ] 3. (depends on 1.) Streaming 2.0
    - a) Add option to `yield`/`yield from` instead of `return` 
    - b) Deprecate the old streamer classes; 
    - c) Add a new class to print the stream into the screen. For beam methods, build a class that prints up to the point where all beams agree with each other.
    - d) thoroughly document and communicate this feature
    - e) enable streaming into the screen with `pipeline`
- [ ] 4. (depends on 2.) Separate stage 1 into a set of functions as described in the diagram. Add unit tests.
- [ ] 5. (depends on 2.) Separate stage 2 into a set of functions as described in the diagram. Add unit tests. Move the preparation of common model inputs here, such as `position_ids`.
- [ ] 6. (depends on 2.) Separate stage 3 into a set of functions as described in the diagram. Add unit tests. Deprecate `LogitsWarper` in this step (it's a copy of `LogitsProcessor`)
- [ ] 7. (depends on 2.) Separate stage 5 into a set of functions as described in the diagram. Add unit tests.
- [ ] 8. Add a new document page walking through the structure of `generate`

[From this point onwards the tasks are only a sketch, need more detailed planning when we get there]
- [ ] 9. Reduce if/elses through templates (e.g. LLMs have a certain default for `prepare_inputs_for_generation`, VLMs also have their special preprocessing steps, ...)
- [ ] 10. Play around with caching of some blocks to determine whether it speeds up generation
- [ ] 11. Rework `prepare_inputs_for_generation` ?
- [ ] 12. Remove `generate` from models that have a custom implementation
- [ ] (other tasks, TBD)

Comments:
Adding the WIP label, so you don't get pinged by the stale bot 🤖 

Could you elaborate on the "prefill" component? My impression is that this step converts the prompt into a KV cache, i.e. "pre-filling" the KV component for the tokens that are fixed. If that's correct, this component could probably serve double duty (as an exit point from the `generate` procedure) for users who just want logprobs/scores for the prompt.

IMHO, upsteam of (part of?) the "generate outputs" step in the decoding loop should be a templated `_prepare_outputs` function whose job is just to attach output attributes to the appropriate output class. Design POC via #29545, concretely: 
https://github.com/coreweave/transformers/blob/dmarx.output_streamer/src/transformers/generation/utils.py#L354-L378


@dmarx 

> Could you elaborate on the "prefill" component? My impression is that this step converts the prompt into a KV cache, i.e. "pre-filling" the KV component for the tokens that are fixed. If that's correct, this component could probably serve double duty (as an exit point from the generate procedure) for users who just want logprobs/scores for the prompt.

Correct, it can be made a public function with that additional purpose. The difference between the prefill for `generate` and obtaining the scores for the prompt is that in the former, we only want to keep the past KV. The different output needs suggest to me that a stand-alone public function is preferable to an alternate exit to `generate` :D Added this to the notes of the prefill stage in the diagram.

> IMHO, upsteam of (part of?) the "generate outputs" step in the decoding loop should be a templated _prepare_outputs function whose job is just to attach output attributes to the appropriate output class. Design POC via https://github.com/huggingface/transformers/pull/29545, concretely:
https://github.com/coreweave/transformers/blob/dmarx.output_streamer/src/transformers/generation/utils.py#L354-L378

That's a good idea, adding it to the diagram!

Thank you for the suggestions 💛 



Hi @gante . I see this issue has been open for 2 months. Do you have any plans to start the 1st step?  If it relies on other tasks and cannot start recently, do you mind considering #31564 as the temporary solution?

@gante maybe you could create sub-issues to track progress on dependent sub tasks and centralize visibility for any in-progress work?

@gante is there any workaround for this in the meantime? Would a custom streamer class be able to pass through the custom model outputs, it seems the streamer class may only be given the token and not other inputs right?

I have the same issue as https://github.com/huggingface/transformers/pull/29545 which brought me here

@tylerweitzman after this plan is complete, streaming should yield all model outputs (and not just the tokens).

The tricky part is timeline, as this competes with other projects 🤗 I'm currently working on #28981 and #32685, which have a higher priority

@tylerweitzman workaround available here: https://github.com/huggingface/transformers/pull/29545