Add TOVA press

[SimJeg]

TOVA press as requested in #2

[maxjeblick]

Thanks a lot for adding TOVA, LGTM!

[hassidm]

Thanks a lot for adding TOVA!

Just a small clarification:
The attn_weights at this line are with shape [batch, heads, seq_len, seq_len] or [batch, heads, window_size, seq_len]? for TOVA in the prefilling phase the former should be implemented (and if I understand correctly, the former is implemented according to link or link).
After computing the attention weights with shape [batch, heads, seq_len, seq_len], the compression should correspond to the last token attention weights (averaged across heads):
torch.mean(attn_weights[:,:,-1:,:],dim=1) (which implemented correctly in th PR :) )

Thanks again, and sorry for the misunderstanding :)

[SimJeg]

After #L58, the shape is [batch, heads, window_size, seq_len] and not [batch, heads, seq_len, seq_len] because we only compute attention for the last window_size tokens.

There is no need to compute the full attention weight matrix that would cost a lot of memory. We only need attn_weights = full_attn_weights[:, :, -window_size:, :].

So in our code,

# attn_weights -> shape [batch, heads, window_size=1, seq_len]
scores = attn_weights.mean(1) # -> shape [batch, window_size=1, seq_len]
scores = scores.repeat(1, keys.shape[1], 1) # -> shape [batch, num_key_value_heads, seq_len]

Do you confirm everything is as you expect it ?

[hassidm]

Thanks for the quick response!

I'm sorry, if I understand correctly this is not exactly the same computation for deeper layers (it is the same for the first attention layer).

For the prefilling case, if token i attended to previous tokens in previous attention layers, the hidden representation will differ from the one where the same token did not attend to other tokens in previous layers. In the current implementation, the non-window tokens (which are all tokens but one) never attend any other tokens, which will result in non contextualized non-window token representations which probably results in (at least) slightly worse performance.

Regarding the memory cost, I agree that this implementation is more costly as we need to store (for a short period) the whole attention matrix of a specific layer, but the KV cache memory stays compressed.

[SimJeg]

The current code is doing exactly the same thing what you first posted in the issue here.

In the current implementation, the non-window tokens (which are all tokens but one) never attend any other tokens

No, the press is applied after the forward pass (in a pytorch hook) and only modify the KV cache, not the hidden states. So whatever you prune in the first layers, the input hidden states of the current layer are not impacted. Hidden states during pre-filling are 100% independent from the press.

I hope this clarifies

[hassidm]

Thank you for clarifying it, and I apologize for the misunderstanding. As I'm not very familiar with this repository, I wanted to ensure that the implemented code functions as intended.

LGTM, and thanks again!