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