H2O implementation

[gopikrishnajha]

In the update_kv function of H2OKVCluster class, I see this code.

attn_weights = torch.matmul(query_states[..., -self.window_size:, :], key_states.transpose(2, 3)) / math.sqrt(head_dim)

As far as I know there is no concept of window in H2O. Shouldn't the entire query_states matrix be considered for attn_weights computation? Why are you only snipping out the window_size part from the query states to be considered for matrix multiplication here?

[Zefan-Cai]

Thank you for pointing out!
The current code is inconsistent with standard H2O as you mentioned. We have tested performance with or without the entire query_state to calculate the attention. We found that this has very limited influence on the performance. We will address this inconsistency in the updated code.