chunked prefill implementation is incorrect

[SmerkyG]

System Info

Latest and older chunked prefill implementations unfortunately contain a significant error, most recently in _prefill:
https://github.com/huggingface/transformers/blob/v5.0.0rc1/src/transformers/generation/utils.py#L3849

position_ids is set incorrectly in the chunked prefill codepath, leading to the wrong RoPE application and bad outputs.

                model_kwargs["position_ids"] = model_kwargs["cache_position"].unsqueeze(0)

This sets the position ids to the chunk being cached, which does not account for all the positions.

Omitting this line fixes the problem, but you may want to initialize the position_ids or even decoder_position_ids for some edge-case use cases. I was unable to find such an initialization that worked and seemed to match the rest of the inits in other parts of the generation utils code.

Who can help?

@gante

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

Run lm eval harness or any eval using a specified prefill_chunk_size, even 999999.

Expected behavior

Gets wrong outputs because the cache positions are specified as being incorrect, leading to wrong RoPE positions being applied.

[SmerkyG]

Update: the difference appears to be between:

               model_inputs[position_ids_key] = model_kwargs["cache_position"].unsqueeze(0)

which is done inside of prefill and

                position_ids = attention_mask.long().cumsum(-1) - 1
                position_ids.masked_fill_(attention_mask == 0, 1)
                model_inputs[position_ids_key] = position_ids[:, past_length:current_length] 

which is used implicitly when the position_ids are calculated automatically after prefill.

The two will mismatch at the end when left padding is employed.

[zucchini-nlp]

Right, we need to prepare position ids based on attention mask if the input is padded. It will work for most cases though there are some models where position ids are prepared differently. I don't think those rare cases will be supported with chunked prefill anyway, so we can assume the general case. Do you want to submit a PR?

Ideally we had a plan to sort out position id preparation and allow users to pass their own positions. It got a bit low in priority though, I will try to revive it this month

[SmerkyG]

I unfortunately don't think I understand the current implementation and position_id constraints enough to offer a PR. To me it seems that position_ids are calculated in two different ways in the code: either as a cumsum of the attention mask or maybe just as an arange in other cases in the code? If so, I think that seems like it needs to be unified to avoid problems. For my current use case I even want to adjust the attention mask during chunked prefill/generation, which will cause problems with the cumsum approach.

The simplest immediate (temporary) bugfix while the issues are considered might be to simply remove the line that reads:
model_kwargs["position_ids"] = model_kwargs["cache_position"].unsqueeze(0)
thereby allowing that cumsum position_ids generation code to be called implicitly. At least that way the same code would be used to generate the position_ids in both chunked prefill and other parts of generation.

[zucchini-nlp]

it needs to be unified to avoid problems.

it definitely does and I had a draft PR long in the past trying to do that. Sadly there were many edge cases with old models where positions are prepared very differently. I want to work on it for sure, though it will take some time until we merge the PR

In the meantime we can patch chunked_prefill to prepare positions in the most common way, i.e. by checking pad positions in the attention mask. yes, simply deleting the line does what the same from what I see

[marcndo]

Hello @zucchini-nlp, i have opened a draft PR linked to this issue. Please have a look when you have a moment.