Swapping attention in a pretrained model for inference

[kabachuha]

Consider we have a LLM, which had been pretrained with quadratic attention, and we want to extend its context size/improve performance. And for this purpose we only swap the attention computation from q,k,v to this rebased linear flash attention.

Similar, still quadratic, attention swap examples include using FlashAttention in XFormers or ScaledDotProduct in Torch2.

Assuming we don't do a backward pass, so no weird gradients breaking the weights. Will the LLM continue inferring more or less fine or will it break down? (Perplexity/loss/qa/needle in stack would be interesting to see)