Cherry-Pick

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

@@ -10,6 +10,8 @@
 from ...ops import *
 
 from .impl import (
+    layernorm_fwd_fl,
+    layernorm_bwd_fl,
     rmsnorm_fwd_fl,
     rmsnorm_bwd_fl,
     multi_tensor_scale_fl,
@@ -160,6 +162,50 @@ def te_general_grouped_gemm(
         )
 
     # Other granular functions
+    def layernorm_fwd(
+        self,
+        input: torch.Tensor,
+        weight: torch.Tensor,
+        bias: Optional[torch.Tensor],
+        eps: float,
+        ln_out: Any,
+        quantizer: Any,
+        otype: DType,
+        sm_margin: int,
+        zero_centered_gamma: bool,
+    ) -> List[Any]:
+        return layernorm_fwd_fl(
+            input=input,
+            weight=weight,
+            bias=bias,
+            eps=eps,
+            ln_out=ln_out,
+            quantizer=quantizer,
+            odtype=otype,
+            sm_margin=sm_margin,
+            zero_centered_gamma=zero_centered_gamma,
+        )
+
+    def layernorm_bwd(
+        self,
+        dz: torch.Tensor,
+        x: torch.Tensor,
+        mu: torch.Tensor,
+        rsigma: torch.Tensor,
+        gamma: torch.Tensor,
+        sm_margin: int,
+        zero_centered_gamma: bool,
+    ) -> List[Any]:
+        return layernorm_bwd_fl(
+            dy=dz,
+            x=x,
+            mu=mu,
+            rsigma=rsigma,
+            gamma=gamma,
+            sm_margin=sm_margin,
+            zero_centered_gamma=zero_centered_gamma,
+        )
+
     def rmsnorm_fwd(
         self,
         input: Any,

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

@@ -7,3 +7,4 @@
 from .fused_adam import *
 from .multi_tensor import *
 from .softmax import *
+from .normalization import *

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

@@ -71,7 +71,6 @@ def generic_gemm_fl(
 
     assert not gelu and gelu_in is None, "Triton-Based General Gemm do not support gelu now"
     assert quantizer is None, "Triton-Based General Gemm do not support quantization now"
-    assert bias is None, "Triton-Based General Gemm do not support bias now"
 
     alpha = validate_gemm_scale(alpha, True)
     beta = validate_gemm_scale(beta, accumulate)
@@ -95,7 +94,18 @@ def generic_gemm_fl(
     A_comp = A.T if transA else A
     B_comp = B.T if transB else B
 
-    out1 = flag_gems.mm(B_comp, A_comp)
+    bias_grad = None
+    if grad:
+        out1 = flag_gems.mm(B_comp, A_comp)
+        if bias is not None:
+            bias_grad = flag_gems.sum_dim(B, dim=[0])
+    else:
+        # NOTE(wqq) flag_gems.addmm uses beta for bias scaling (Y = alpha * WX + beta * bias),
+        # unlike the beta here (for scaling D). Always set to 1.
+        if bias is not None:
+            out1 = flag_gems.addmm(bias, B_comp, A_comp, beta=1, alpha=alpha)
+        else:
+            out1 = flag_gems.mm(B_comp, A_comp)
 
     if shape_b_changed:
         out1 = out1.view(s, b, -1)
@@ -104,7 +114,6 @@ def generic_gemm_fl(
     if torch_out_dtype is not None and out1.dtype != torch_out_dtype:
         out1 = out1.to(torch_out_dtype)
 
-    bias_grad = None
     gelu_input = None
     extra_output_ret = None
 

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

@@ -0,0 +1,62 @@
+# Copyright (c) 2025, BAAI. All rights reserved.
+#
+# See LICENSE for license information.
+
+import torch
+import flag_gems
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+
+def layernorm_fwd_fl(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: Optional[torch.Tensor],
+    eps: float,
+    ln_out: Any,
+    quantizer: Any,
+    odtype: Any,
+    sm_margin: int,
+    zero_centered_gamma: bool,
+) -> List[Any]:
+    if zero_centered_gamma:
+        # weight_adj = 1 + weight
+        weight_adj = flag_gems.add(1, weight)
+    else:
+        weight_adj = weight
+
+    y, mean, rstdevs = flag_gems.layer_norm(
+        input,
+        [input.shape[-1]],
+        weight_adj,
+        bias=bias,
+        eps=eps,
+    )
+
+    if rstdevs.shape != input.shape[:-1]:
+        rstdevs = rstdevs.view(input.shape[:-1])
+
+    return y, mean, rstdevs
+
+
+def layernorm_bwd_fl(
+    dy: torch.Tensor,
+    x: torch.Tensor,
+    mu: torch.Tensor,
+    rsigma: torch.Tensor,
+    gamma: torch.Tensor,
+    sm_margin: int,
+    zero_centered_gamma: bool,
+) -> List[Any]:
+    # When zero_centered_gamma is True, forward uses (1 + gamma) as weight
+    # So backward needs to use (1 + gamma) for computing dx
+    if zero_centered_gamma:
+        gamma_adj = flag_gems.add(1, gamma)
+    else:
+        gamma_adj = gamma
+
+    dummy_bias = torch.zeros(x.shape[-1], dtype=x.dtype, device=x.device)
+    dx, dw, db = flag_gems.layer_norm_backward(
+        dy, x, None, mu, rsigma, weight=gamma_adj, bias=dummy_bias, output_mask=[True, True, True]
+    )
+
+    return dx, dw, db

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

@@ -42,6 +42,22 @@ def register_builtins(registry) -> None:
     is_avail = backend.is_available
 
     impls = [
+        OpImpl(
+            op_name="layernorm_fwd",
+            impl_id="default.flagos",
+            kind=BackendImplKind.DEFAULT,
+            fn=_bind_is_available(backend.layernorm_fwd, is_avail),
+            vendor=None,
+            priority=150,
+        ),
+        OpImpl(
+            op_name="layernorm_bwd",
+            impl_id="default.flagos",
+            kind=BackendImplKind.DEFAULT,
+            fn=_bind_is_available(backend.layernorm_bwd, is_avail),
+            vendor=None,
+            priority=150,
+        ),
         OpImpl(
             op_name="rmsnorm_fwd",
             impl_id="default.flagos",

transformer_engine/plugin/core/backends/vendor/kunlunxin/kunlunxin.py

@@ -357,6 +357,21 @@ def multi_tensor_compute_scale_and_scale_inv(
         epsilon: float,
     ) -> None:
         tex = self._get_tex()
-        return self.multi_tensor_compute_scale_and_scale_inv(
+        return tex.multi_tensor_compute_scale_and_scale_inv(
             chunk_size, noop_flag, tensor_lists, max_fp8, force_pow_2_scales, epsilon
         )
+
+    def multi_tensor_compute_scale_inv_e8m0(
+        self,
+        chunk_size: int,
+        noop_flag: torch.Tensor,
+        tensor_lists: List[List[torch.Tensor]],
+        block_len: int,
+    ) -> None:
+        tex = self._get_tex()
+        return tex.multi_tensor_compute_scale_inv_e8m0(
+            chunk_size,
+            noop_flag,
+            tensor_lists,
+            block_len,
+        )

transformer_engine/plugin/core/backends/vendor/kunlunxin/register_ops.py

@@ -173,6 +173,14 @@ def register_builtins(registry) -> None:
             vendor="KUNLUNXIN",
             priority=100,
         ),
+        OpImpl(
+            op_name="multi_tensor_compute_scale_inv_e8m0",
+            impl_id="vendor.kunlunxin",
+            kind=BackendImplKind.VENDOR,
+            fn=_bind_is_available(backend.multi_tensor_compute_scale_inv_e8m0, is_avail),
+            vendor="KUNLUNXIN",
+            priority=100,
+        ),
     ]
 
     registry.register_many(impls)

transformer_engine/plugin/tests/run_all_tests.py

@@ -10,6 +10,8 @@
 from test_softmax import SoftmaxTests
 from test_optimizer import OptimizerTests
 from test_flash_attention import FlashAttentionTests
+from test_te_general_grouped import grouped_gemmTests
+from test_policy import run_all_tests
 
 
 def main():
@@ -27,6 +29,7 @@ def main():
         SoftmaxTests(device=device),
         OptimizerTests(device=device),
         FlashAttentionTests(device=device),
+        grouped_gemmTests(device=device),
     ]
 
     results = []
@@ -49,6 +52,8 @@ def main():
     print(f"Total: {total_passed}/{total_tests} test suites passed")
     print("=" * 70)
 
+    run_all_tests()
+
     return 0 if all(success for _, success in results) else 1