apply flagos te_groups_gemm op (#55)

- add ```te_general_grouped_gemm``` op for flagos backend, base on
flag_gems
- support both forward and backward computation, distinguished by
```grad```

Author: chai-xiaonan

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
+
+
+# This function can represent both forward and backward computations.
+# When grad is False (forward computation), the 'bias' is bias;
+# When grad is True (backward computation/gradient calculation), the 'bias' is grad_bias;
+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 or grad_bias
+    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))
+
+    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:
+                    flag_gems.copy_(D[i], flag_gems.zeros(D[i].shape))
+            else:
+                out = flag_gems.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:
+                flag_gems.copy_(bias[i], flag_gems.zeros(bias[i].shape))
+            if (
+                len(pre_gelu_out) > i
+                and pre_gelu_out[i] is not None
+                and pre_gelu_out[i].numel() != 0
+            ):
+                flag_gems.copy_(pre_gelu_out[i], flag_gems.zeros(pre_gelu_out[i].shape))
+            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:
+                flag_gems.copy_(pre_gelu_out[i], out)
+                out = flag_gems.gelu(out)
+        else:
+            out = flag_gems.mm(a, b)
+
+            # Apply dGELU epilogue if requested
+            if has_pre_gelu:
+                out = flag_gems.gelu_backward(out, pre_gelu_out[i])
+
+            # Compute bias gradients if requested
+            if has_bias:
+                bias_grad = flag_gems.sum_dim(out, dim=[0])
+                if accumulate:
+                    flag_gems.add_(bias[i], bias_grad)
+                else:
+                    flag_gems.copy_(bias[i], bias_grad)
+
+        if not single_output:
+            # Store output
+            if accumulate:
+                flag_gems.add_(D[i], out.to(D[i].dtype))
+            else:
+                flag_gems.copy_(D[i], out.to(D[i].dtype))
+        else:
+            temp_D.append(out.to(D[0].dtype))
+
+    if single_output:
+        if temp_D:
+            temp = flag_gems.cat(temp_D, dim=0)
+            flag_gems.copy_(D[0], 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",

transformer_engine/plugin/tests/test_te_general_grouped.py

@@ -0,0 +1,169 @@
+import torch
+
+from transformer_engine.plugin.test_utils import (
+    get_available_backends,
+    get_backend,
+    TestCase,
+    generate_random_tensor,
+)
+
+
+class grouped_gemmTests(TestCase):
+    def __init__(self, device="cpu"):
+        super().__init__(
+            "Moe permute Operations",
+            "Test correctness of all moe permute operations across backends",
+        )
+        self.backends = get_available_backends()
+        self.device = device
+
+    def test_grouped_gemm_equivalence(self, grad, has_bias, has_pre_gelu, single_output):
+        print(
+            "\n test te_general_grouped_gemm"
+            f" grad:{grad} has_bias:{has_bias},has_pre_gelu:{has_pre_gelu},single_output:{single_output}"
+        )
+        import transformer_engine_torch_nv as tex
+
+        num_gemms = 2
+        m, k, n = 128, 32, 64
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        dtype = torch.float16
+
+        if dtype == torch.float16:
+            te_dtype = tex.DType.kFloat16
+        elif dtype == torch.float32:
+            te_dtype = tex.DType.kFloat32
+        elif dtype == torch.bfloat16:
+            te_dtype = tex.DType.kBFloat16
+        else:
+            raise ValueError(f"不支持的 dtype: {torch_dtype}")
+
+        torch.manual_seed(42)
+
+        A_list = [torch.randn((k, n), device=device, dtype=dtype) for _ in range(num_gemms)]
+        B_list = [torch.randn((m, k), device=device, dtype=dtype) for _ in range(num_gemms)]
+
+        bias_list_py_bias = [
+            (
+                torch.randn(n, device=device, dtype=dtype)
+                if has_bias
+                else torch.empty(0, device=device, dtype=dtype)
+            )
+            for _ in range(num_gemms)
+        ]
+        bias_list_te = [b.clone() for b in bias_list_py_bias]
+
+        pre_gelu_list_py = [
+            (
+                torch.randn(m, n, device=device, dtype=dtype)
+                if has_pre_gelu
+                else torch.empty(0, device=device, dtype=dtype)
+            )
+            for _ in range(num_gemms)
+        ]
+        pre_gelu_list_te = [p.clone() for p in pre_gelu_list_py]
+
+        if single_output:
+            D_list_py = [torch.empty(m * num_gemms, n, device=device, dtype=dtype)]
+            D_list_te = [torch.empty(m * num_gemms, n, device=device, dtype=dtype)]
+        else:
+            D_list_py = [torch.empty(m, n, device=device, dtype=dtype) for _ in range(num_gemms)]
+            D_list_te = [torch.empty(m, n, device=device, dtype=dtype) for _ in range(num_gemms)]
+        workspace_py = [torch.empty(1024 * 1024, device=device, dtype=torch.uint8)]
+        workspace_te = [torch.empty(1024 * 1024, device=device, dtype=torch.uint8)]
+
+        tex.te_general_grouped_gemm(
+            A_list,
+            False,
+            B_list,
+            False,
+            D_list_te,
+            te_dtype,
+            [],
+            bias_list_te,
+            te_dtype,
+            single_output,
+            pre_gelu_list_te,
+            grad,
+            workspace_te,
+            1024 * 1024,
+            False,
+            False,
+            0,
+        )
+
+        for backend_name in self.backends:
+            backend = get_backend(backend_name)
+            print("backend:", backend)
+            try:
+                bias_list_py = [b.clone() for b in bias_list_py_bias]
+                backend.te_general_grouped_gemm(
+                    A_list,
+                    False,
+                    B_list,
+                    False,
+                    D_list_py,
+                    te_dtype,
+                    [],
+                    bias_list_py,
+                    te_dtype,
+                    single_output,
+                    pre_gelu_list_py,
+                    grad,
+                    workspace_py,
+                    1024 * 1024,
+                    False,
+                    False,
+                    0,
+                )
+
+                for py_d, te_d in zip(D_list_py, D_list_te):
+                    self.assert_close(
+                        py_d, te_d, rtol=1e-3, atol=1e-3, msg="Output D tensors mismatch!"
+                    )
+
+                if not grad and has_pre_gelu:
+                    for py_p, te_p in zip(pre_gelu_list_py, pre_gelu_list_te):
+                        self.assert_close(
+                            py_p, te_p, rtol=1e-3, atol=1e-3, msg="Pre-GELU out tensors mismatch!"
+                        )
+
+                if grad or has_bias:
+                    for py_b, te_b in zip(bias_list_py, bias_list_te):
+                        self.assert_close(
+                            py_b, te_b, rtol=1e-3, atol=1e-3, msg="Bias gradient tensors mismatch!"
+                        )
+                print(f"    ✓ {backend_name}")
+            except NotImplementedError:
+                self.skipped += 1
+                print(f"    ⊘  {backend_name} (not implemented)")
+            except Exception as e:
+                self.failed += 1
+                print(f"    ✗ Test failed: {e}")
+
+    def run_all_tests(self):
+        print("\n" + "=" * 60)
+        print("=" * 60)
+        print(f"Available backends: {', '.join(self.backends)}")
+
+        # gemm tests
+        self.test_grouped_gemm_equivalence(False, False, False, False)
+        self.test_grouped_gemm_equivalence(False, True, False, False)
+        self.test_grouped_gemm_equivalence(False, False, True, False)
+
+        self.test_grouped_gemm_equivalence(False, False, False, True)
+        self.test_grouped_gemm_equivalence(False, True, False, True)
+        self.test_grouped_gemm_equivalence(False, False, True, True)
+        return self.report()
+
+
+def main():
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"Using device: {device}")
+    test_suite = grouped_gemmTests(device=device)
+    success = test_suite.run_all_tests()
+    return 0 if success else 1
+
+
+if __name__ == "__main__":
+    exit(main())