Use fused implementation of RoPE in MultiHeadAttention (#658)

* Use fused implementation of RoPE in MultiHeadAttention
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Fix freqs dtype
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

---------
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

tests/pytorch/fused_attn/test_fused_attn.py

@@ -811,7 +811,7 @@ def _run_transformer_layer(
     rotary_pos_emb = None
     if RoPE:
         PE = RotaryPositionEmbedding(dim=config.head_dim)
-        rotary_pos_emb = PE(config.max_seqlen_q).to(dtype=dtype, device="cuda")
+        rotary_pos_emb = PE(config.max_seqlen_q).to(device="cuda")
 
     # Set up model
     block = (

transformer_engine/pytorch/attention.py

@@ -3625,8 +3625,8 @@ class MultiheadAttention(torch.nn.Module):
         # apply relative positional encoding (rotary embedding)
         if rotary_pos_emb is not None:
             q_pos_emb, k_pos_emb = rotary_pos_emb
-            query_layer = apply_rotary_pos_emb(query_layer, q_pos_emb, self.qkv_format)
-            key_layer = apply_rotary_pos_emb(key_layer, k_pos_emb, self.qkv_format)
+            query_layer = apply_rotary_pos_emb(query_layer, q_pos_emb, self.qkv_format, fused=True)
+            key_layer = apply_rotary_pos_emb(key_layer, k_pos_emb, self.qkv_format, fused=True)
 
         context_layer = self.core_attention(
             query_layer,