KV Cache getting prefilled and emptied during training phase

[adityavipradas]

Although the KV cache implementation for inference looks good, the block_kv_cache is getting prefilled in the training phase leading to increase in compute and memory consumption.

nanoVLM/models/vision_language_model.py

Line 51 in 098db57

logits, _ = self.decoder(combined_embd, attention_mask=attention_mask) # Not logits yet, but easier to return like this

[ariG23498]

This is a really good catch! Do you want to contribute this feature?

I think a flag with training=True would help not create the cache at all.

[adityavipradas]

Yes sure. I will be happy to update it.