Added HybridCache class and function

Signed-off-by: Ann <[email protected]>

Author: quic-akuruvil

QEfficient/transformers/cache_utils.py

@@ -9,7 +9,7 @@
 from typing import Any, Dict, Optional, Tuple
 
 import torch
-from transformers.cache_utils import DynamicCache, EncoderDecoderCache
+from transformers.cache_utils import DynamicCache, EncoderDecoderCache, HybridCache
 
 from QEfficient.customop import (
     CtxGatherFunc,
@@ -428,3 +428,50 @@ def from_legacy_cache(
                     cache.cross_attention_cache.update(key_states, value_states, layer_idx)
                     cache.is_updated[layer_idx] = True
         return cache
+
+
+class QEffHybridCache(HybridCache):
+    def update(
+        self,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        layer_idx: int,
+        cache_kwargs: Optional[Dict[str, Any]] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        position_ids = cache_kwargs.get("position_ids")
+        sliding_window_pattern = cache_kwargs.get("sliding_window_pattern")
+        is_sliding_layer = torch.tensor(bool((layer_idx + 1) % sliding_window_pattern))
+        N = self.key_cache[layer_idx].shape[2]
+
+        kv_position_ids = torch.where((~is_sliding_layer | (position_ids == -1)), position_ids, position_ids % (N - 1))
+
+        kv_position_ids = torch.where(
+            is_sliding_layer & (position_ids.max() >= (N - 1) * 2), (position_ids + 1) % N, kv_position_ids
+        )
+
+        valid_mask = (kv_position_ids != -1).unsqueeze(1).unsqueeze(-1)
+        key_states = torch.where(valid_mask == 1, key_states, torch.zeros_like(key_states))
+        value_states = torch.where(valid_mask == 1, value_states, torch.zeros_like(value_states))
+        self.key_cache[layer_idx] = CtxScatterFunc.apply(self.key_cache[layer_idx], kv_position_ids, key_states)
+        self.value_cache[layer_idx] = CtxScatterFunc.apply(self.value_cache[layer_idx], kv_position_ids, value_states)
+        k_out, v_out = self.key_cache[layer_idx], self.value_cache[layer_idx]
+
+        # Original Gather
+        ctx_len = self.key_cache[layer_idx].shape[2]
+        ctx_indices = torch.arange(ctx_len)[None, None, ...]
+        gather_limit = kv_position_ids.max(1, keepdim=True).values.unsqueeze(1)
+        invalid_mask = ctx_indices > gather_limit
+        if torch.onnx.is_in_onnx_export():
+            invalid_idx_value = torch.iinfo(torch.int32).max
+        else:
+            invalid_idx_value = 0
+        ctx_indices = torch.where(invalid_mask, invalid_idx_value, ctx_indices)
+
+        all_indices = torch.arange(N) + kv_position_ids.max() + 1
+        rolling_indices = torch.where(all_indices > N - 1, all_indices % N, all_indices)
+        final_indices = torch.where((is_sliding_layer & (position_ids.max() >= (N - 1))), rolling_indices, ctx_indices)
+        k_out = CtxGatherFunc.apply(k_out, final_indices)
+        v_out = CtxGatherFunc.apply(v_out, final_indices)
+        ctx_v_out = torch.where(invalid_mask.unsqueeze(-1), torch.tensor(0.0, dtype=torch.float32), v_out)
+        v_out = torch.where((is_sliding_layer & (position_ids.max() >= (N - 1))), v_out, ctx_v_out)
+        return k_out, v_out