support assisted decoding in ipex 2.4

[jiqing-feng]

Support assisted decoding if ipex >= 2.5

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[jiqing-feng]

Hi @echarlaix . I have fixed all your comments, please re-run the tests. Thx!

[jiqing-feng]

Hi @IlyasMoutawwakil , could you please review this PR? Thx!

[IlyasMoutawwakil]

Can you please add an explanation in the PR of why this patching is necessary ?
Do assistant models interfere with torchscript ?

[jiqing-feng]

Can you please add an explanation in the PR of why this patching is necessary ? Do assistant models interfere with torchscript ?

Yes. In assisted decoding, both the target model and the assistant model need to call crop_pask_key_values. The patched IPEXModel has a different format of past_key_values, so we need our own crop functions to apply IPEX KV-cache.

Do you want me to write it as a comment in the codes?

[IlyasMoutawwakil]

Yes. In assisted decoding, both the target model and the assistant model need to call crop_pask_key_values. The patched IPEXModel has a different format of past_key_values, so we need our own crop functions to apply IPEX KV-cache.
Do you want me to write it as a comment in the codes?

Thanks for the explanation ! One last code change and it's good for me 🤗

[jiqing-feng]

Yes. In assisted decoding, both the target model and the assistant model need to call crop_pask_key_values. The patched IPEXModel has a different format of past_key_values, so we need our own crop functions to apply IPEX KV-cache.
Do you want me to write it as a comment in the codes?

Thanks for the explanation ! One last code change and it's good for me 🤗

I have fixed your required changes. The _unpatch_crop_past_key_values() can be safely added at the end of the generation because it will revert the function crop_pask_key_values if we patch and keep the same if we didn't patch.

BTW, please help to trigger the CI if it's okay for you, thx!

[jiqing-feng]

Hi @echarlaix . I have enabled assisted decoding tests in ipex 2.4. Please take a review, thx~

[echarlaix]

could this be unpatch after export ?

[jiqing-feng]

As we discussed in Teams, this is the only way to enable all assisted decoding cases:

1. transformers target model + ipex draft model
2. ipex target model + transformers draft model
3. ipex target model + ipex draft model

The _crop_past_key_values function is the same level as the model, we cannot do un-patch inside the generate function because it will run after generate, see here.

I have checked the model type inside the _ipex_crop_past_key_values. It only has impact on IPEX model, transformers model will go into the original function, so there is no risk even we don't unpatch.

[jiqing-feng]

Hi @echarlaix . I can do the unpatch for target model only, please look the new changes