Clarification about GroupQueryAttention behavior

[ZackyLake]

It's about com.microsoft.GroupQueryAttention

The document says:

past_key (optional) : T_CACHE
past state key with support for format BNSH. When past_key uses same tensor as present_key(k-v cache), it is of length max_sequence_length... otherwise of length past_sequence_length.

But when compile the model, there's no information about what tensors being bound to past/present kv.
Does it mean we can use the shape to infer runtime tensor binding? Eg, when past_kv and present_kv is of same shape, app must bind the same tensor?

Also, consider an edge case, that past_kv and present_kv are all max_seq_len == past_seq_len, what should happen when there's no room to append new kv?

We are working on EP development, so would like to know the exact behavior of GQA.
We want to avoid copy if needed, and better to know during model compilation.

[Mayuresh Varerkar (MayureshV1)]

Tianlei Wu (@tianleiwu).. Can you please help answer the above question from ZackyLake.

[Tianlei Wu (tianleiwu)]

You can use shape info to assist compiling model (like static shape for shared buffer). In most cases in ort-genai, shared buffer is used for past and present kv, and we can also add a runtime check to validate tensor sharing (points to same address) for security.

For the case that there is no room to append new kv. The operator shall fail (ORT_THROW).

[ZackyLake]

Thanks for the reply!
I was thinking about the case that past/present are static shape of 128 and 256, with new_token_len=128, but seems that's not allowed.

Can I summarize as follows:

1. If past and present are static shape, they must be equal.

2a. If past and present are same shape(static or dynamic with same dynamic-dim), they must be bond to same tensor.

2b. If past and present are different shape(dynamic shape), present should be past + new and they should be bond to different tensors?

[ZackyLake]

Tianlei Wu (@tianleiwu) latest gqa spec has changes for SWA, so there's 2 questions:

The spec says "The op then keeps only the min(total_sequence_length, cache_capacity) most recent tokens, contiguously".
Does it means when total_seq_len is long, the present_kv should be full of past tokens in every iteration? That means every iteration needs to do an eviction, causing a copy of size being kvcache capacity.

For EP enabling, can we choose to do "conditional and greedy eviction", which means "everytime when putting new tokens causing overflow, evict at many token as possible only keep latest W"?
That could heavily reduce the need of kvcache copy and improve performance, but not sure if it's acceptable by spec.

attention_bias is shape: (batch_size or 1, num_heads or 1, sequence_length, total_sequence_length).

So does total_seq_len mean past+cur ? But static model requires a static shape for input and past+cur is dynamic.

If it's actually max_seq_len which is the same with past/present kv, then in SWA is it actually cache_capacity, which means it's not a fully mask, but trimmed using similar rule with kvcache?