What does the update_cache function in MultiHeadAttention do?

[fedorovgv]

I am currently working with the MultiHeadAttention class and found the update_cache function. As far as I understand it does nothing at this moment and is a template for the future, am I right? and if it's true, can you explain what this function will do?

NeMo/nemo/collections/asr/parts/submodules/multi_head_attention.py

Lines 154 to 165 in 9f94649

	def update_cache(self, key, value, query, cache, cache_next):
	if cache is not None:
	q_length = query.size(1)
	q_input = query
	key = value = torch.cat((cache[self._cache_id], key), dim=1)
	if cache_next is not None:
	q_keep_size = torch.tensor(q_length - self.cache_drop_size, dtype=torch.int64).clip(min=1)
	keep_in_cache_next(cache=cache, cache_next=cache_next, keep_size=q_keep_size, cache_id=self._cache_id)
	cache_next[self._cache_id, :, -q_keep_size:, :] = q_input[:, :q_keep_size, :]
	return key, value, query

[titu1994]

It's used in chunk aware transformers. cc @VahidooX

[VahidooX]

The caching is being used during the inference when cache-aware streaming Conformer is being used. During the training, it is skipped.

[fedorovgv]

Thanks a lot!