Untied Ulysses

Intermediate buffers are too large! (all Q,K,V, x2 because of all-to-all)

Let's leverage the multi-head aspect of attention:

Inp-All-to-All

Attention

Out-All-to-All

HBM

QKV

All-to-All

OOM