Add FlagOS Triton fused RoPE kernels

transformer_engine/plugin/core/backends/flagos/flagos.py

@@ -22,6 +22,10 @@
     scaled_masked_softmax_forward_fl,
     scaled_masked_softmax_backward_fl,
     te_general_grouped_gemm_fl,
+    fused_rope_forward_fl,
+    fused_rope_backward_fl,
+    fused_qkv_rope_forward_fl,
+    fused_qkv_rope_backward_fl,
 )
 
 
@@ -352,6 +356,101 @@ def multi_tensor_adam_param_remainder(
             weight_decay,
         )
 
+    # fused apply rope
+    def fused_rope_forward(
+        self,
+        input: torch.Tensor,
+        freqs: torch.Tensor,
+        start_positions: Optional[torch.Tensor],
+        qkv_format: NVTE_QKV_Format,
+        interleaved: bool,
+        cu_seqlens: Optional[torch.Tensor],
+        cp_size: int,
+        cp_rank: int,
+    ) -> torch.Tensor:
+        return fused_rope_forward_fl(
+            input,
+            freqs,
+            start_positions,
+            qkv_format,
+            interleaved,
+            cu_seqlens,
+            cp_size,
+            cp_rank,
+        )
+
+    def fused_rope_backward(
+        self,
+        output_grads: torch.Tensor,
+        freqs: torch.Tensor,
+        start_positions: Optional[torch.Tensor],
+        qkv_format: NVTE_QKV_Format,
+        interleaved: bool,
+        cu_seqlens: Optional[torch.Tensor],
+        cp_size: int,
+        cp_rank: int,
+    ) -> torch.Tensor:
+        return fused_rope_backward_fl(
+            output_grads,
+            freqs,
+            start_positions,
+            qkv_format,
+            interleaved,
+            cu_seqlens,
+            cp_size,
+            cp_rank,
+        )
+
+    def fused_qkv_rope_forward(
+        self,
+        qkv_input: torch.Tensor,
+        q_freqs: torch.Tensor,
+        k_freqs: torch.Tensor,
+        start_positions: Optional[torch.Tensor],
+        qkv_split_arg_list: List[int],
+        qkv_format: NVTE_QKV_Format,
+        interleaved: bool,
+        cp_size: int,
+        cp_rank: int,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        return fused_qkv_rope_forward_fl(
+            qkv_input,
+            q_freqs,
+            k_freqs,
+            start_positions,
+            qkv_split_arg_list,
+            qkv_format,
+            interleaved,
+            cp_size,
+            cp_rank,
+        )
+
+    def fused_qkv_rope_backward(
+        self,
+        q_grad_out: torch.Tensor,
+        k_grad_out: torch.Tensor,
+        v_grad_out: torch.Tensor,
+        q_freqs: torch.Tensor,
+        k_freqs: torch.Tensor,
+        qkv_split_arg_list: List[int],
+        qkv_format: NVTE_QKV_Format,
+        interleaved: bool,
+        cp_size: int,
+        cp_rank: int,
+    ) -> torch.Tensor:
+        return fused_qkv_rope_backward_fl(
+            q_grad_out,
+            k_grad_out,
+            v_grad_out,
+            q_freqs,
+            k_freqs,
+            qkv_split_arg_list,
+            qkv_format,
+            interleaved,
+            cp_size,
+            cp_rank,
+        )
+
     # Misc
     def get_cublasLt_version(self) -> int:
         return 110000

transformer_engine/plugin/core/backends/flagos/impl/__init__.py

@@ -8,3 +8,4 @@
 from .multi_tensor import *
 from .softmax import *
 from .normalization import *
+from .trition import *

transformer_engine/plugin/core/backends/flagos/impl/trition/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) 2025, BAAI. All rights reserved.
+#
+# See LICENSE for license information.
+
+from .fused_rope import *