Yeah, the problem is indeed in the attention computation.
You can do something like that but it's far from optimal.
From memory consumption perspective, the right way to do it, is to never materialize the intermediate matrices.
You can do it, by using a customop, that compute
att = scaledAttention(Q,K,V) and the gradient dQ,dK,dV = scaledAttentionBackward(Q,K,V,att,datt)
The memory needed for these ops is the memory to store Q,K,V,attn,dQ,dK,dV,dattn + extra temporary memory.
When you do the work to minimize memory consumption, this extra temporary memory is really small : 6attention_horizon^2number_of_core_running_in_parallel numbers.
But even though there is not much re computation, this kernel won't run as fast due to the pattern of memory access, unless you spend some time manually optimizing it.
The place to do it is at the level of the autodiff framework aka tensorflow or pytorch, with low level c++/cuda code.
Anybody can write some custom kernel, but deploying, maintaining them and distributing them is a nightmare. So the only people that could and should have done it, are the tensorflow or pytorch guys.
In fact they probably have, but it's considered a strategic advantage and reserved for internal use only.
The mere mortals like us, have to use some workarounds (splitting matrices, Kheops, gradient checkpointing... ) to not be too much penalized by the limited ops of the out of the box autodiff frameworks like tensorflow or torch.