Possible issue with the semantic correctness of KV caches.

[eomiso]

Possible issue with Semantic correctness of KV caches.

The shape of KVcache for mobileVLM:

The KV Cache size : [batch, num_kv_heads, past_seq_len, head_dim]

Current CoreML, iOS code initializes KVcache with size [1, 16, 1, 128] and fills it with zeros(src/mobilevlm/iOS/mobilevlm/mobilevlm/Core/Manager/bufferManager.swift:48). Then when the prefill happens, the shape would become [1, 16, 1 + seq_len, 128]. The prefill starts with the dummy kvcache passed on from the swift code.
If the dummy kvcache is required by the iOS, then the attention_mask should make sure that the dummy kvcache is unattended.

current attention_mask: [1, 1, 1, 1, ...] (len: curEmbed.shape[1].intValue + 1 ). This should be changed to [0, 1, 1, .... 1].

[eomiso]

I might be missing some details of the implementations so take this as a naive friendly review.

It seems like the pytorch allows shapes like [1, 16, 0, 128], while swift doesn't. I wonder how pytorch handles this.

[johyeongseob]

Hi eomiso. Thank you for your Interest!

1. The right attention mask is [0, 1, 1, ... , 1] in the prefill stage. I agree! But the zero kv cache consists with 'zero', so there is no effect if fisrt element (of slot) of attention mask is 0 or 1. I mean both work well. :)

2. In the prefill stage, Pytorch can handle both [24, 2, 1, 16, 0, 128] and [24, 2, 1, 16, 1, 128] as a kv cache. (Even thought omission can be allowed!) But the CoreML cannot allow the zero slot. So I add a dummy slot in the kv cache.

If you have any additional question, please let me know!

Thank you and Regards.

Hyeongseob Jo

[eomiso]

The attentions masks are converted to a vector with elements either [0, -inf] called attention mask bias. (This is done inside the model, meaning there is a component in the compute graph that does this).

scores = Q@K^T
scores = scores + attention_mask_bias
weights = softmax(scores)
output = weights @ V

For example:

raw scores: [0.0, 2.1, 0.5, 1.2]
attention masks: [1, 1, 1, ,1]
softmax: [0.07, 0.58, 0.12, 0.23]

May have a marginal effect but if you don't mask out the dummy input using the attention mask the dummy input is still attended regardless of them being zero-ed out.