Transformers compatability

[maxjeblick]

PR description

Updates kvpress to upcoming transformers version 4.56.0.

Follwoing issues/PRs need to be fixed before merging this PR:

• QuantoQuantizedCache is not working huggingface/transformers#40099
• Wrong function signature introduced in PR #40029 breaking flash attention huggingface/transformers#40097
• [Flash Attention] Fix flash attention integration huggingface/transformers#40002

PR was tested on huggingface/transformers#40002

TODOS:

• rerun (subset) of ruler benchmark to ensure no regression

Checklist

• Tests are working (make test)
• Code is formatted correctly (make style, on errors try fix with make format)
• Copyright header is included
• All commits are signed-off using git commit -s
• (new press) mypress_press.py is in the presses directory
• (new press) MyPress is in __init__.py
• (new press) README.md is updated with a 1 liner about the new press in the Available presses section
• (new press) New press is in the default_presses list in tests/default_presses.py
• (new press) A docstring is provided that follows the same structure as the existing ones

[maxjeblick]

To test with current transformers PR being open:
uv pip install git+https://github.com/vasqu/transformers.git@fix-fa-integration

[manueldeprada]

Also note that now we compute the flash attention varlen kwargs early in generate:
https://github.com/vasqu/transformers/blob/cb89dbee4694ca1fec3d5733c296b729291bb507/src/transformers/generation/utils.py#L681-L692

to avoid expensive recomputations on each attention pass. This is the case since huggingface/transformers#39474 and huggingface/transformers#40002

Unfortunately, this means kvpress might cause an out of bound access, since the condition assert keys.shape[0] == cu_seqlens_k[-1] must hold and it doesnt if kvpress changes the keys and the computation of cu_seqlens_k happened earlier, in generate.

Have a look at 9986c31 as a possible fix. This script tests for such failure:

import os
os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
from transformers import AutoModelForCausalLM, AutoTokenizer
from kvpress import KnormPress

ckpt = "meta-llama/Meta-Llama-3-8B-Instruct"
device = "cuda"

model = AutoModelForCausalLM.from_pretrained(ckpt, torch_dtype="auto").to(device)
model.set_attn_implementation("flash_attention_2")
tok = AutoTokenizer.from_pretrained(ckpt)
inputs = tok("Hello, how are you? bla bla how are you? this is some text lala ddd", return_tensors="pt").to(device)

with KnormPress(0.8)(model):
    outputs = model.generate(**inputs, max_new_tokens=10, do_sample=False)

[maxjeblick]

Thanks a lot for the comments @manueldeprada, this is really appreciated!

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[alessiodevoto]

/ok to test f37b768

[alessiodevoto]

Thanks a lot @maxjeblick for handling this, and @Jack-Yu-815 for the review!