transformer_engine/pytorch/attention.py · 209b8e5a801a7747adfa42235aedeb6bfd805874 · OpenDAS / TransformerEngine

[PyTorch] Relax the contiguous check for flash attention (#1176) · 0ee5ccda

Xin Yao authored Sep 19, 2024



* relax contiguous check for flash attention
Signed-off-by: Xin Yao <xiny@nvidia.com>

* force contiguous for cp
Signed-off-by: Xin Yao <xiny@nvidia.com>

---------
Signed-off-by: Xin Yao <xiny@nvidia.com>

0ee5ccda

attention.py 372 KB

Replace attention.py