apply flagos te_groups_gemm op

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

@@ -19,6 +19,7 @@
     generic_gemm_fl,
     scaled_masked_softmax_forward_fl,
     scaled_masked_softmax_backward_fl,
+    te_general_grouped_gemm_fl,
 )
 
 
@@ -118,6 +119,46 @@ def generic_gemm(
             beta,
         )
 
+    def te_general_grouped_gemm(
+        self,
+        A: List[Any],
+        transa: bool,
+        B: List[Any],
+        transb: bool,
+        D: Optional[List[torch.Tensor]],
+        D_type: DType,
+        m_splits: List[int],
+        bias: List[torch.Tensor],
+        bias_type: DType,
+        single_output: bool,
+        pre_gelu_out: List[torch.Tensor],
+        grad: bool,
+        workspace: List[torch.Tensor],
+        workspaceSizes: int,
+        accumulate: bool,
+        use_split_accumulator: bool,
+        math_sm_count: int,
+    ) -> Optional[List[torch.Tensor]]:
+        return te_general_grouped_gemm_fl(
+            A,
+            transa,
+            B,
+            transb,
+            D,
+            D_type,
+            m_splits,
+            bias,
+            bias_type,
+            single_output,
+            pre_gelu_out,
+            grad,
+            workspace,
+            workspaceSizes,
+            accumulate,
+            use_split_accumulator,
+            math_sm_count,
+        )
+
     # Other granular functions
     def rmsnorm_fwd(
         self,

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

@@ -10,6 +10,7 @@
 
 __all__ = [
     "generic_gemm_fl",
+    "te_general_grouped_gemm_fl",
 ]
 
 _DTYPE_TO_TORCH = {
@@ -115,3 +116,107 @@ def generic_gemm_fl(
         return D, bias_grad, gelu_input, extra_output_ret
     else:
         return out1, bias_grad, gelu_input, extra_output_ret
+
+
+def te_general_grouped_gemm_fl(
+    B: List[torch.Tensor],
+    transb: bool,
+    A: List[torch.Tensor],
+    transa: bool,
+    D: Optional[List[torch.Tensor]],
+    D_type: Any,
+    m_splits: List[int],
+    bias: List[torch.Tensor],
+    bias_type: Any,
+    single_output: bool,
+    pre_gelu_out: List[torch.Tensor],
+    grad: bool,
+    workspace: List[torch.Tensor],
+    workspaceSize: int,
+    accumulate: bool,
+    use_split_accumulator: bool,
+    math_sm_count: int,
+) -> Optional[List[torch.Tensor]]:
+    if single_output and D is None:
+        raise ValueError("not implemented, D should be allocated for single output case.")
+
+    num_gemms = len(A)
+    if D is None:
+        D = []
+        for i in range(num_gemms):
+            m = A[i].shape[1] if transa else A[i].shape[0]
+            n = B[i].shape[0] if transb else B[i].shape[1]
+            D.append(torch.empty((m, n), dtype=D[i].dtype, device=A[0].device))
+
+    def gelu_backward(grad_output, x):
+        # Approximation of GELU derivative commonly used in Transformer Engine
+        cdf = 0.5 * (1.0 + flag_gems.erf(x / flag_gems.sqrt(2.0)))
+        pdf = flag_gems.exp(-0.5 * x * x) / flag_gems.sqrt(2.0 * math.pi)
+        return flag_gems.mul(grad_output, (cdf + x * pdf))
+
+    temp_D = []
+    for i in range(num_gemms):
+        # Handle the special case of zero-element inputs
+        if A[i].numel() == 0 or B[i].numel() == 0:
+            if not single_output:
+                if D[i].numel() != 0 and not accumulate:
+                    D[i].zero_()
+            else:
+                out = torch.zeros(A[i].shape[0], B[i].shape[1])
+            if grad and len(bias) > i and bias[i] is not None and bias[i].numel() != 0:
+                bias[i].zero_()
+            if (
+                len(pre_gelu_out) > i
+                and pre_gelu_out[i] is not None
+                and pre_gelu_out[i].numel() != 0
+            ):
+                pre_gelu_out[i].zero_()
+            continue
+        a = A[i].t() if transa else A[i]
+        b = B[i].t() if transb else B[i]
+        # Determine presence of epilogue tensors
+        has_bias = len(bias) > i and bias[i] is not None and bias[i].numel() > 0
+        has_pre_gelu = (
+            len(pre_gelu_out) > i and pre_gelu_out[i] is not None and pre_gelu_out[i].numel() > 0
+        )
+
+        # Forward Pass calculation
+        if not grad:
+            if has_bias:
+                # Fused matrix multiplication and bias addition
+                out = flag_gems.addmm(bias[i], a, b)
+            else:
+                out = flag_gems.mm(a, b)
+
+            # Apply GELU epilogue if pre_gelu_out is provided
+            if has_pre_gelu:
+                pre_gelu_out[i].copy_(out)
+                out = flag_gems.gelu(out)
+        else:
+            out = flag_gems.mm(a, b)
+            if has_pre_gelu:
+                out = gelu_backward(out, pre_gelu_out[i])
+
+            # Compute bias gradients if requested
+            if has_bias:
+                bias_grad = out.sum(dim=0)
+                if accumulate:
+                    bias[i].add_(bias_grad)
+                else:
+                    bias[i].copy_(bias_grad)
+
+        if not single_output:
+            # Store output
+            if accumulate:
+                D[i].add_(out.to(D[i].dtype))
+            else:
+                D[i].copy_(out.to(D[i].dtype))
+        else:
+            temp_D.append(out.to(D[0].dtype))
+
+    if single_output:
+        if temp_D:
+            temp = torch.cat(temp_D, dim=0)
+            D[0].copy_(temp)
+
+    return bias

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

@@ -66,6 +66,14 @@ def register_builtins(registry) -> None:
             vendor=None,
             priority=150,
         ),
+        OpImpl(
+            op_name="te_general_grouped_gemm",
+            impl_id="default.flagos",
+            kind=BackendImplKind.DEFAULT,
+            fn=_bind_is_available(backend.te_general_grouped_gemm, is_avail),
+            vendor=None,
+            priority=150,
+        ),
         OpImpl(
             op_name="multi_tensor_scale",
             impl_id="default.flagos",