Add scaled_masked_softmax_forward/backward for flagos backend (flagos-ai#52)

Add two functions for flagos backend, based on flaggems

- scaled_masked_softmax_forward
- scaled_masked_softmax_backend

Author: lxd-cumt

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

@@ -17,6 +17,8 @@
     multi_tensor_adam_param_remainder_fl,
     multi_tensor_l2_norm_fl,
     generic_gemm_fl,
+    scaled_masked_softmax_forward_fl,
+    scaled_masked_softmax_backward_fl,
 )
 
 
@@ -160,6 +162,23 @@ def rmsnorm_bwd(
     def get_fused_attn_backend(self, *args, **kwargs) -> int:
         return NVTE_Fused_Attn_Backend.NVTE_No_Backend
 
+    # Softmax functions
+    def scaled_masked_softmax_forward(
+        self,
+        input: torch.Tensor,
+        mask: torch.Tensor,
+        scale_factor: Union[float, torch.Tensor],
+    ) -> torch.Tensor:
+        return scaled_masked_softmax_forward_fl(input, mask, scale_factor)
+
+    def scaled_masked_softmax_backward(
+        self,
+        output_grad_: torch.Tensor,
+        softmax_results_: torch.Tensor,
+        scale_factor: float,
+    ) -> torch.Tensor:
+        return scaled_masked_softmax_backward_fl(output_grad_, softmax_results_, scale_factor)
+
     # multi-tensor functions
     def multi_tensor_scale(
         self,

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

@@ -6,3 +6,4 @@
 from .rmsnorm import *
 from .fused_adam import *
 from .multi_tensor import *
+from .softmax import *

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

@@ -0,0 +1,61 @@
+import torch
+from typing import Union
+import flag_gems
+
+
+def scaled_masked_softmax_forward_fl(
+    input: torch.Tensor,
+    mask: torch.Tensor,
+    scale_factor: Union[float, torch.Tensor],
+) -> torch.Tensor:
+    # Ensure `mask` and 'scale_factor' is on the same device as `input`.
+    if mask.device != input.device:
+        mask = flag_gems.to_copy(mask, device=input.device)
+    if isinstance(scale_factor, torch.Tensor):
+        if scale_factor.device != input.device:
+            scale_factor = flag_gems.to_copy(scale_factor, device=input.device)
+
+    # Keep semantics aligned with TE CUDA scaled_masked_softmax:
+    # - integer/bool mask: masked iff mask == 1, masked logits set to -10000.0
+    # - float mask: treated as additive bias in logit space
+    if mask.dim() == 4 and mask.size(1) == 1 and input.dim() == 4:
+        mask = mask.expand_as(input)
+
+    scaled = flag_gems.mul(input, scale_factor)
+    if mask.is_floating_point():
+        mask_f = flag_gems.to_copy(mask, device=input.device, dtype=scaled.dtype)
+        scaled = flag_gems.add(scaled, mask_f)
+        return flag_gems.softmax(scaled, dim=-1)
+
+    # Avoid using `mask == 1` (torch op) since on some devices it may fall back to CPU,
+    # which would break Triton kernels inside flag_gems.
+    cond = flag_gems.eq_scalar(mask, 1)
+    scaled = flag_gems.masked_fill(scaled, cond, -10000.0)
+    all_masked = flag_gems.all_dim(cond, dim=-1, keepdim=True)
+    out = flag_gems.softmax(scaled, dim=-1)
+    return flag_gems.masked_fill(out, all_masked, 0.0)
+
+
+def scaled_masked_softmax_backward_fl(
+    output_grad_: torch.Tensor,
+    softmax_results_: torch.Tensor,
+    scale_factor: float,
+) -> torch.Tensor:
+    orig_dtype = output_grad_.dtype
+    # Compute in float32 for numerical stability.
+    output_grad_f32 = flag_gems.to_copy(output_grad_, dtype=torch.float32)
+    softmax_output_f32 = flag_gems.to_copy(
+        softmax_results_, dtype=torch.float32, device=output_grad_.device
+    )
+    if isinstance(scale_factor, torch.Tensor):
+        if scale_factor.device != output_grad_.device:
+            scale_factor = flag_gems.to_copy(scale_factor, device=output_grad_.device)
+
+    # term = softmax_output_f32 * output_grad_f32
+    term = flag_gems.mul(softmax_output_f32, output_grad_f32)
+    # sum_term = sum(term, dim=-1, keepdim=True)
+    sum_term = flag_gems.sum_dim(term, dim=[-1], keepdim=True)
+    # grad_softmax = softmax_output_f32 * (output_grad_f32 - sum_term)
+    grad_softmax = flag_gems.mul(softmax_output_f32, flag_gems.sub(output_grad_f32, sum_term))
+    grad_scaled = flag_gems.mul(grad_softmax, scale_factor)
+    return flag_gems.to_copy(grad_scaled, dtype=orig_dtype)

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

@@ -132,6 +132,22 @@ def register_builtins(registry) -> None:
             vendor=None,
             priority=150,
         ),
+        OpImpl(
+            op_name="scaled_masked_softmax_forward",
+            impl_id="default.flagos",
+            kind=BackendImplKind.DEFAULT,
+            fn=_bind_is_available(backend.scaled_masked_softmax_forward, is_avail),
+            vendor=None,
+            priority=150,
+        ),
+        OpImpl(
+            op_name="scaled_masked_softmax_backward",
+            impl_id="default.flagos",
+            kind=BackendImplKind.DEFAULT,
+            fn=_bind_is_available(backend.scaled_masked_softmax_backward, is_avail),
+            vendor=None,
+            priority=150,
+        ),
         OpImpl(
             op_name="get_cudnn_version",
             impl_id="default.flagos",