transformers

6446f1fc - Fix offloaded cache device mismatch on hybrid models (#46748)

Fix offloaded cache device mismatch on hybrid models (#46748) * Fix offloaded cache device mismatch on hybrid models Generating with cache_implementation="offloaded" crashes on hybrid models that interleave attention and linear-attention (Mamba/SSM) layers, with a cpu/accelerator device mismatch when concatenating the KV cache. Only attention layers go through the offloading update path; the prefetch of the next layer kept landing on the interleaved linear-attention layers (a no-op) and never brought the next attention layers KV back to the accelerator, so after prefill every attention layers KV was stranded on CPU. Make prefetch skip layers that are not offloaded and target the next attention layer that actually goes through the offloading update path, wrapping around if needed. This is a no-op for non-hybrid caches. Signed-off-by: Ting Sun <suntcrick@gmail.com> * fix for hybrid models * fix for hybrid --------- Signed-off-by: Ting Sun <suntcrick@gmail.com> Co-authored-by: Cyril Vallez <cyril.vallez@gmail.com>