enable StaticCache for assisted generation
Author: yao-matrix
Pull request: yes

@gante , I implemented a version for this issue: https://github.com/huggingface/transformers/issues/32946. Pls help comment, and I can iterate, thx.


Comments:
@gante , could you pls take a look？ thx

> @yao-matrix hey, gante is currently on a long vacation so I reviewed the PR for him. Thanks for adding support for this, Super cool work!
> 
> I left a few comments and also we'll need tests in `tests/generation/test_utils.py` file. I guess static cache now works with all types of candidate generators right?

@zucchini-nlp , test_utils CI pass rate is the same before and after this PR, as below. So no regressions are introduced.
before:
=========================== short test summary info ============================
FAILED tests/generation/test_utils.py::GenerationIntegrationTests::test_assisted_decoding_encoder_decoder_shared_encoder
FAILED tests/generation/test_utils.py::GenerationIntegrationTests::test_assisted_decoding_num_assistant_tokens_heuristic_schedule
FAILED tests/generation/test_utils.py::GenerationIntegrationTests::test_assisted_generation_early_exit
FAILED tests/generation/test_utils.py::GenerationIntegrationTests::test_custom_logits_processor
FAILED tests/generation/test_utils.py::GenerationIntegrationTests::test_default_max_length_warning
FAILED tests/generation/test_utils.py::GenerationIntegrationTests::test_eos_token_id_int_and_list_beam_search
FAILED tests/generation/test_utils.py::GenerationIntegrationTests::test_eos_token_id_int_and_list_top_k_top_sampling
FAILED tests/generation/test_utils.py::GenerationIntegrationTests::test_generate_compile_fullgraph_tiny
FAILED tests/generation/test_utils.py::GenerationIntegrationTests::test_generated_length_assisted_generation
FAILED tests/generation/test_utils.py::GenerationIntegrationTests::test_max_new_tokens_encoder_decoder
FAILED tests/generation/test_utils.py::GenerationIntegrationTests::test_min_length_if_input_embeds
FAILED tests/generation/test_utils.py::GenerationIntegrationTests::test_model_kwarg_assisted_decoding_decoder_only
FAILED tests/generation/test_utils.py::GenerationIntegrationTests::test_model_kwarg_assisted_decoding_encoder_decoder
FAILED tests/generation/test_utils.py::GenerationIntegrationTests::test_model_kwarg_encoder_signature_filtering
FAILED tests/generation/test_utils.py::GenerationIntegrationTests::test_prepare_inputs_for_generation_decoder_llm
FAILED tests/generation/test_utils.py::GenerationIntegrationTests::test_speculative_decoding_equals_regular_decoding
FAILED tests/generation/test_utils.py::GenerationIntegrationTests::test_stop_sequence_stopping_criteria
====== 17 failed, 51 passed, 19 skipped, 13 warnings in 133.78s (0:02:13) ======

after:
=========================== short test summary info ============================
FAILED tests/generation/test_utils.py::GenerationIntegrationTests::test_assisted_decoding_encoder_decoder_shared_encoder
FAILED tests/generation/test_utils.py::GenerationIntegrationTests::test_assisted_decoding_num_assistant_tokens_heuristic_schedule
FAILED tests/generation/test_utils.py::GenerationIntegrationTests::test_assisted_generation_early_exit
FAILED tests/generation/test_utils.py::GenerationIntegrationTests::test_custom_logits_processor
FAILED tests/generation/test_utils.py::GenerationIntegrationTests::test_default_max_length_warning
FAILED tests/generation/test_utils.py::GenerationIntegrationTests::test_eos_token_id_int_and_list_beam_search
FAILED tests/generation/test_utils.py::GenerationIntegrationTests::test_eos_token_id_int_and_list_top_k_top_sampling
FAILED tests/generation/test_utils.py::GenerationIntegrationTests::test_generate_compile_fullgraph_tiny
FAILED tests/generation/test_utils.py::GenerationIntegrationTests::test_generated_length_assisted_generation
FAILED tests/generation/test_utils.py::GenerationIntegrationTests::test_max_new_tokens_encoder_decoder
FAILED tests/generation/test_utils.py::GenerationIntegrationTests::test_min_length_if_input_embeds
FAILED tests/generation/test_utils.py::GenerationIntegrationTests::test_model_kwarg_assisted_decoding_decoder_only
FAILED tests/generation/test_utils.py::GenerationIntegrationTests::test_model_kwarg_assisted_decoding_encoder_decoder
FAILED tests/generation/test_utils.py::GenerationIntegrationTests::test_model_kwarg_encoder_signature_filtering
FAILED tests/generation/test_utils.py::GenerationIntegrationTests::test_prepare_inputs_for_generation_decoder_llm
FAILED tests/generation/test_utils.py::GenerationIntegrationTests::test_speculative_decoding_equals_regular_decoding
FAILED tests/generation/test_utils.py::GenerationIntegrationTests::test_stop_sequence_stopping_criteria
====== 17 failed, 51 passed, 19 skipped, 13 warnings in 133.78s (0:02:13) ======

> LGTM! We need some tests and then I am requesting review from the core maintainer, after that we can merge

thx for reviewing.

@yao-matrix no worries is some tests are failing and are not related to PR changes. Might be just flaky or will be fixed on `main` by us. From what I see the only CI test affected by PR is this one + need to see if new test passes for all models

`tests/models/gemma2/test_modeling_gemma2.py::Gemma2ModelTest::test_assisted_decoding_with_num_logits_to_keep`

@zucchini-nlp , any more comments for me to iterate? Thx.

@yao-matrix no, the only thing is the CI which is failing now. I showed the relevant test in prev comment and if you can add one more test in `tests/generation/test_utils.py` which would test static cache with assisted generation. That is all actually

At the end you need to run `make style` to pass CI check on codestyle. Feel free to tag the core maintainer @ ArthurZucker for review when tests are ready and CI is green or tag me if you need help/have questions :)  

@zucchini-nlp  & @ArthurZucker , I cleaned the PR, pls help review. I plan to enable SlidingWindowCache and HybridCache as next step, thx.


> Looks very nice, but we need to add a `compile` test to make sure this is compile compatible! The whole point of static cache is -> compile! 🤗

@ArthurZucker  i added a `test_assisted_decoding_compile` case based on `test_generate_compile`, `forward_only` test pass for llama, `end_to_end` test fail for the same reason as Joao commented in `test_generate_compile`. 

@ArthurZucker @zucchini-nlp , pls let me know any further comments, thx. BTW, checked the failed ci case, not relevant to my changes.

Thanks, re-triggered the tests, let's wait for the core maintainer 

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_34797). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@ArthurZucker , @zucchini-nlp , I am thinking is it possible we leave this PR in 2024, :).

@zucchini-nlp @ArthurZucker , any further comments on this?

@yao-matrix Arthur Zucker might be a bit busy, so tagging @gante to review is possible. We'll be able to merge after we get one more review

> @yao-matrix Arthur Zucker might be a bit busy, so tagging @gante to review is possible. We'll be able to merge after we get one more review

thx, @gante , could you pls help review? Thx.

@yao-matrix any plans to continue this effort? Compiling the assistant is an important use case IMO and it only really makes sense with a static cache. 

Hi, I am get the issue "ValueError: assisted generate is not supported with Static cache classes`", when try speculative decoding for Gemma family. 

`model = AutoModelForCausalLM.from_pretrained(
            args.target_model,
            torch_dtype=torch.bfloat16,
            trust_remote_code=True,
            attn_implementation="sdpa" if 'gemma' in args.target_model else "flash_attention_2", 
            device_map={"": device},  # Map all modules to the specified device
        )`

`draft_model = AutoModelForCausalLM.from_pretrained(
                args.draft_model,
                torch_dtype=torch.bfloat16,
                trust_remote_code=True,
                attn_implementation="sdpa" if 'gemma' in args.target_model else "flash_attention_2", 
                device_map={"": device},  # Map all modules to the specified device
            )`

`output_ids = model.generate(
                        input_ids,
                        attention_mask=attn_mask,
                        max_new_tokens=args.max_tokens,
                        do_sample=do_sample,
                        use_cache=True,
                        assistant_model=draft_model,
                        output_hidden_states=True,
                        temperature=args.temperature if args.temperature > 0 else 1.0,
                        num_beams=1,
                    )`