SIGN IN SIGN UP

implement position encoding for shifted tokens

Pull Request resolved: https://github.com/pytorch/executorch/pull/6646

In AttentionSink, it uses tokens' positions in the KVCache instead of the actual text. When tokens get shifted in KVCache, it needs to update q and k's position embedding.

In the original [implementation](https://github.com/mit-han-lab/streaming-llm) of AttentionSink with Rope, it caches the original q and k in KVCache and apply position embedding during inference.

This PR adds `RopeWithAttentionSink`. It assumes that q and k are already encoded with their original position. When we shift tokens, we reapply the position delta. This has two benefits:
- minimize our code since our existing `llama_transformer` applies rope embedding before doing KVCache update
- avoid performance regression when tokens are not shifted because we don't need to reapply position encoding in KVCache for them
ghstack-source-id: 255579838

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

---------

Co-authored-by: Lunwen He <lwhecser@gmail.com>
P
pytorchbot committed
6b738410e400b173fdda78bea352aa5eb334e751
Parent: 27638c3
Committed by GitHub <noreply@github.com> on 11/27/2024, 8:03:49 PM