Add gemma 2

[ArthurZucker]

What does this PR do?

Adds support for gemma2

[LysandreJik]

Good! Only need to update to the 2 in the integration tests

[LysandreJik]

Ok!

[turboderp]

Please excuse this if I'm just not reading the code correctly (?), but I'm struggling to understand the intended function of the hybrid cache. Here, in modeling_gemma2.py as well as the two other attention functions, the past keys/values are updated like so:

        if past_key_value is not None:
            # sin and cos are specific to RoPE models; cache_position needed for the static cache
            cache_kwargs = {
                "sin": sin,
                "cos": cos,
                "sliding_window": self.sliding_window,
                "cache_position": cache_position,
            }
            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)

While the function definition for HybridCache.update looks like:

    def update(
        self,
        key_states: torch.Tensor,
        value_states: torch.Tensor,
        layer_idx: int,
        cache_kwargs: Optional[Dict[str, Any]] = None,
        sliding_window: Optional[int] = None,
    ) -> Tuple[torch.Tensor]:
        cache_position = cache_kwargs.get("cache_position")
        self.key_cache[layer_idx] = self.key_cache[layer_idx].to(device=key_states.device)
        self.value_cache[layer_idx] = self.value_cache[layer_idx].to(device=value_states.device)
        k_out = self.key_cache[layer_idx]
        v_out = self.value_cache[layer_idx]
        if sliding_window:
            update_fn = self._sliding_update
        else:
            update_fn = self._static_update
    ...

It doesn't read the sliding_window argument from the kwargs, so the default value is always used and the _sliding_update function is never selected, even on layers that use a sliding_window. Is this right?

[ArthurZucker]

that is a good catch and not intended, we will update this and fix it. The typo is that we should either pass the sliding windo directly, or get the sliding window form the cache kwargs. I think this stems from a will to make it compile compatible, and a type!