[WIP] drafting a fix - cropping the kv cache

[keyboardAnt]

It seems like we haven't been cropping the KV cache. This PR adds the missing functionality and hopefully addresses the issue raised in huggingface#35029 (comment)

[jmamou]

Contrarily to UAG, I don't think we need to set remove_from_pkv > 0 in USD (like in regular AG/SD) since draft/target tokens are aligned

[keyboardAnt]

Contrarily to UAG, I don't think we need to set remove_from_pkv > 0 in USD (like in regular AG/SD) since draft/target tokens are aligned

How does the KV cache get updated after rejecting draft tokens?

[jmamou]

Contrarily to UAG, I don't think we need to set remove_from_pkv > 0 in USD (like in regular AG/SD) since draft/target tokens are aligned

How does the KV cache get updated after rejecting draft tokens?

new_cache_size = input_ids.shape[-1] - 1
self.assistant_kwargs["past_key_values"] = _crop_past_key_values(
    self.assistant_model, self.assistant_kwargs["past_key_values"], new_cache_size - 1
)

Here, input_ids corresponds to assistant_input_ids with validated token only + additional token(s) from the target validation

[gauravjain14]

@keyboardAnt - Is this change still applicable if #7 is merged?

[jmamou]

@keyboardAnt - Is this change still applicable if #7 is merged?

@gauravjain14 no