BUG: New models like Gemma 3 and Command A start losing coherency

[vlbosch]

Gemma 3 loses coherency quite soon when generating text. The issue is not with MLX, see ml-explore/mlx-lm#58

For Cohere Command A, it starts later on than with Gemma 3, but that models also starts producing gibberish, see ml-explore/mlx-lm#11 (comment)

[vlbosch]

I debugged the issue and can confirm that it has to do with the caching. The current KVCache is a default dynamic cache that grows dynamically with an increasing number of keys and values. Gemma 3 and Command A implement sliding window attention with a varying attention window and pattern.

I got both models working after implementing a form of SWA-caching. I will see if I can clean it up and submit a PR for this.

[nath1295]

Hi @vlbosch thank you for sticking with the package, a bit caught up at work lately so I have little time to work on the project.

However, I do have a few modifications in mind:

1. Incorporating mask for better batch inference
2. Handling the cache better for multimodal models and different types of cache for newer model architectures

It seems that you have encounter an issue due to the current way of handling cache. Can you find any existing code of SWA-caching from the mlx-lm that we can integrate or at least use as a reference point?

The challenge with sliding window is the slicing of cache after we find the appropriate cache history, we might need to redesign that mechanism to store and search multiple caches for these models. Same for cache for images as well. The current approach rely on the length of the cache and the stored list of tokens, which might not be enough for these new cache types.

[nath1295]

In the latest version of the package, I have tested gemma3 models and it seems coherent after a long conversation. The vision is working fine as well. @vlbosch can you verify this by updating the package?