Update on "add attention_sink.py"

This PR adds `KVCacheWithAttentionSink`, which is required for `AttentionSink`. It keeps the first `sink_size` tokens as attention sinks and maintains a sliding window with `window_size` for new tokens.

Note: I am trying to implement and verify `AttentionSink` in eager mode first. So the current implementation may still have some lower errors or performance issue. For example, it does not support the case when dynamic shape is disabled. Will leave these problems to resolve when we are ready to deploy `AttentionSink` to edge.

Differential Revision: [D65235798](https://our.internmc.facebook.com/intern/diff/D65235798/)

[ghstack-poisoned]

Author: helunwencser

examples/models/llama/source_transformation/attention_sink.py

@@ -79,25 +79,25 @@ def update(
             num_to_evict = min(start_pos + seq_length - self.cache_size, seq_length)
 
             # Shift the existing entries to the left
-            # pyre-ignore: Incompatible parameter type [6]
             k_to_keep = self.k_cache.narrow(
                 dim_to_slice,
-                self.sink_size + num_to_evict,
-                self.window_size - num_to_evict,
+                self.sink_size + num_to_evict,  # pyre-ignore [6]
+                self.window_size - num_to_evict,  # pyre-ignore [6]
             ).clone()
-            # pyre-ignore: Incompatible parameter type [6]
             v_to_keep = self.v_cache.narrow(
                 dim_to_slice,
-                self.sink_size + num_to_evict,
-                self.window_size - num_to_evict,
+                self.sink_size + num_to_evict,  # pyre-ignore [6]
+                self.window_size - num_to_evict,  # pyre-ignore [6]
             ).clone()
-            # pyre-ignore: Incompatible parameter type [6]
             k_new_position = self.k_cache.narrow(
-                dim_to_slice, self.sink_size, self.window_size - num_to_evict
+                dim_to_slice,
+                self.sink_size,
+                self.window_size - num_to_evict,  # pyre-ignore [6]
             )
-            # pyre-ignore: Incompatible parameter type [6]
             v_new_position = self.v_cache.narrow(
-                dim_to_slice, self.sink_size, self.window_size - num_to_evict
+                dim_to_slice,
+                self.sink_size,
+                self.window_size - num_to_evict,  # pyre-ignore [6]
             )
             k_new_position.copy_(k_to_keep)
             v_new_position.copy_(v_to_keep)