transformer_engine/pytorch/attention.py · da30634a6c9ccdbb6c587b6c93b1860e4b038204 · OpenDAS / TransformerEngine

[PyTorch] Refactor caching of cumulative sequence lengths (#630) · da30634a

Tim Moon authored Feb 05, 2024



Do not cache sequence lengths based on layer number
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

da30634a

attention.py 162 KB

Replace attention.py