tests/pytorch/test_fused_rope.py · 838345eba4fdd2a169dd9e087d39c30a360e684a · OpenDAS / TransformerEngine

[PyTorch] Let Fused RoPE support CP with THD format (#1238) · 55dcbb4b

Xin Yao authored Oct 12, 2024



* Let Fused RoPE support THD with CP
Signed-off-by: Xin Yao <xiny@nvidia.com>

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

---------
Signed-off-by: Xin Yao <xiny@nvidia.com>
Co-authored-by: Xiaowei Ren <103958965+xrennvidia@users.noreply.github.com>

55dcbb4b

test_fused_rope.py 6.51 KB

Replace test_fused_rope.py