Differential Attention is incompatible with Infini-Attention

[Vectorrent]

You can reproduce the errors like this:

python run.py --dev --no_dashboard --memory --differential

The main conflict is coming from the fact that differential heads are essentially doubled in size, so that they can be split in half, later. Infini-Attention was not designed to handle this, and thus fails here.

I tried to fix, but could not find an elegant solution, which left the math intact. Would love some help here.