Integrate KunLunXin TE-FL backend patches

Move the temporary XTE TE-FL patch behavior into TE-FL native backend
implementations. Register KunLunXin layernorm and GEMM operators, route
attention backend selection through transformer_engine_klx_torch, and add
reference GLU/DGLU fallback implementations.

Author: root

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

@@ -10,6 +10,7 @@
 from .activation import (
     gelu_torch,
     geglu_torch,
+    glu_torch,
     qgelu_torch,
     qgeglu_torch,
     relu_torch,
@@ -21,6 +22,7 @@
     clamped_swiglu_torch,
     dgelu_torch,
     dgeglu_torch,
+    dglu_torch,
     dqgelu_torch,
     dqgeglu_torch,
     drelu_torch,
@@ -71,6 +73,7 @@
     "layernorm_bwd_torch",
     "gelu_torch",
     "geglu_torch",
+    "glu_torch",
     "qgelu_torch",
     "qgeglu_torch",
     "relu_torch",
@@ -82,6 +85,7 @@
     "clamped_swiglu_torch",
     "dgelu_torch",
     "dgeglu_torch",
+    "dglu_torch",
     "dqgelu_torch",
     "dqgeglu_torch",
     "drelu_torch",

transformer_engine/plugin/core/backends/reference/impl/activation.py

@@ -46,6 +46,11 @@ def geglu_torch(input: torch.Tensor, quantizer: Any) -> torch.Tensor:
     return F.gelu(a, approximate="tanh") * b
 
 
+def glu_torch(input: torch.Tensor, quantizer: Any) -> torch.Tensor:
+    a, b = input.chunk(2, dim=-1)
+    return torch.sigmoid(a) * b
+
+
 def qgelu_torch(input: torch.Tensor, quantizer: Any) -> torch.Tensor:
     return input * torch.sigmoid(1.702 * input)
 
@@ -123,6 +128,18 @@ def dgeglu_torch(grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) ->
     return torch.cat([a.grad, b.grad], dim=-1)
 
 
+def dglu_torch(grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> torch.Tensor:
+    a, b = fwd_input.chunk(2, dim=-1)
+    a = a.detach().requires_grad_(True)
+    b = b.detach().requires_grad_(True)
+
+    with torch.enable_grad():
+        y = torch.sigmoid(a) * b
+        y.backward(grad)
+
+    return torch.cat([a.grad, b.grad], dim=-1)
+
+
 def dqgelu_torch(grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> torch.Tensor:
     x = fwd_input.detach().requires_grad_(True)
     with torch.enable_grad():

transformer_engine/plugin/core/backends/reference/reference.py

@@ -15,6 +15,7 @@
     layernorm_bwd_torch,
     gelu_torch,
     geglu_torch,
+    glu_torch,
     qgelu_torch,
     qgeglu_torch,
     relu_torch,
@@ -26,6 +27,7 @@
     clamped_swiglu_torch,
     dgelu_torch,
     dgeglu_torch,
+    dglu_torch,
     dqgelu_torch,
     dqgeglu_torch,
     drelu_torch,
@@ -161,6 +163,9 @@ def gelu(self, input: torch.Tensor, quantizer: Any) -> Any:
     def geglu(self, input: torch.Tensor, quantizer: Any) -> Any:
         return geglu_torch(input, quantizer)
 
+    def glu(self, input: torch.Tensor, quantizer: Any) -> Any:
+        return glu_torch(input, quantizer)
+
     def qgelu(self, input: torch.Tensor, quantizer: Any) -> Any:
         return qgelu_torch(input, quantizer)
 
@@ -203,6 +208,9 @@ def dgelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) ->
     def dgeglu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any:
         return dgeglu_torch(grad, fwd_input, quantizer)
 
+    def dglu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any:
+        return dglu_torch(grad, fwd_input, quantizer)
+
     def dqgelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any:
         return dqgelu_torch(grad, fwd_input, quantizer)
 

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

@@ -101,6 +101,14 @@ def register_builtins(registry) -> None:
             vendor=None,
             priority=50,
         ),
+        OpImpl(
+            op_name="glu",
+            impl_id="reference.torch",
+            kind=BackendImplKind.REFERENCE,
+            fn=_bind_is_available(backend.glu, is_avail),
+            vendor=None,
+            priority=50,
+        ),
         OpImpl(
             op_name="qgelu",
             impl_id="reference.torch",
@@ -190,6 +198,14 @@ def register_builtins(registry) -> None:
             vendor=None,
             priority=50,
         ),
+        OpImpl(
+            op_name="dglu",
+            impl_id="reference.torch",
+            kind=BackendImplKind.REFERENCE,
+            fn=_bind_is_available(backend.dglu, is_avail),
+            vendor=None,
+            priority=50,
+        ),
         OpImpl(
             op_name="dqgelu",
             impl_id="reference.torch",

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

@@ -101,6 +101,37 @@ def rmsnorm_bwd(
         tex = self._get_tex()
         return tex.rmsnorm_bwd(dz, x, rsigma, gamma, sm_margin, zero_centered_gamma)
 
+    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]:
+        tex = self._get_tex()
+        otype = tex.DType(int(otype)) if otype is not None else None
+        return tex.layernorm_fwd(
+            input, weight, bias, eps, ln_out, quantizer, otype, sm_margin, 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]:
+        tex = self._get_tex()
+        return tex.layernorm_bwd(dz, x, mu, rsigma, gamma, sm_margin, zero_centered_gamma)
+
     def multi_tensor_adam(
         self,
         chunk_size: int,
@@ -298,17 +329,16 @@ def get_cudnn_version(self) -> int:
         return 0
 
     def get_attention_backend(self, attention_params=None):
-        from transformer_engine_klx.pytorch import attention
+        tex = self._get_tex()
 
         (
             use_flash_attention,
+            flash_attention_backend,
             use_fused_attention,
             fused_attention_backend,
             use_unfused_attention,
             available_backends,
-        ) = attention.get_attention_backend(attention_params)
-
-        flash_attention_backend = None
+        ) = tex.get_attention_backend(attention_params)
 
         return (
             use_flash_attention,
@@ -375,3 +405,63 @@ def multi_tensor_compute_scale_inv_e8m0(
             tensor_lists,
             block_len,
         )
+
+    def generic_gemm(
+        self,
+        A: Any,
+        transA: bool,
+        B: Any,
+        transB: bool,
+        D: Any,
+        quantizer: Any,
+        output_dtype: Optional[DType],
+        bias: Optional[torch.Tensor],
+        bias_type: DType,
+        gelu: bool,
+        gelu_in: Optional[torch.Tensor],
+        grad: bool,
+        workspace: torch.Tensor,
+        workspace_size: int,
+        accumulate: bool,
+        use_split_accumulator: bool,
+        comm_overlap: Optional[Any] = None,
+        comm_type: Optional[CommOverlapType] = None,
+        extra_output: Optional[torch.Tensor] = None,
+        bulk_overlap: bool = False,
+        alpha: float = 1.0,
+        beta: Optional[float] = None,
+    ) -> List[Any]:
+        tex = self._get_tex()
+        return tex.generic_gemm(
+            A, transA, B, transB, D, quantizer, output_dtype, bias, bias_type,
+            gelu, gelu_in, grad, workspace, workspace_size, accumulate,
+            use_split_accumulator, comm_overlap, comm_type, extra_output,
+            bulk_overlap, alpha, beta,
+        )
+
+    def te_general_grouped_gemm(
+        self,
+        A: List[torch.Tensor],
+        transa: bool,
+        B: List[torch.Tensor],
+        transb: bool,
+        D: Optional[List[torch.Tensor]],
+        D_type,
+        m_splits: List[int],
+        bias: List[torch.Tensor],
+        bias_type,
+        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]]:
+        tex = self._get_tex()
+        return tex.te_general_grouped_gemm(
+            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,
+        )

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

@@ -61,6 +61,22 @@ def register_builtins(registry) -> None:
             vendor="KUNLUNXIN",
             priority=100,
         ),
+        OpImpl(
+            op_name="layernorm_fwd",
+            impl_id="vendor.kunlunxin",
+            kind=BackendImplKind.VENDOR,
+            fn=_bind_is_available(backend.layernorm_fwd, is_avail),
+            vendor="KUNLUNXIN",
+            priority=200,
+        ),
+        OpImpl(
+            op_name="layernorm_bwd",
+            impl_id="vendor.kunlunxin",
+            kind=BackendImplKind.VENDOR,
+            fn=_bind_is_available(backend.layernorm_bwd, is_avail),
+            vendor="KUNLUNXIN",
+            priority=200,
+        ),
         OpImpl(
             op_name="multi_tensor_adam",
             impl_id="vendor.kunlunxin",
@@ -181,6 +197,23 @@ def register_builtins(registry) -> None:
             vendor="KUNLUNXIN",
             priority=100,
         ),
+        # GEMM (XPU via hydrax in transformer_engine_klx_torch.gemm)
+        OpImpl(
+            op_name="generic_gemm",
+            impl_id="vendor.kunlunxin",
+            kind=BackendImplKind.VENDOR,
+            fn=_bind_is_available(backend.generic_gemm, is_avail),
+            vendor="KUNLUNXIN",
+            priority=200,
+        ),
+        OpImpl(
+            op_name="te_general_grouped_gemm",
+            impl_id="vendor.kunlunxin",
+            kind=BackendImplKind.VENDOR,
+            fn=_bind_is_available(backend.te_general_grouped_gemm, is_avail),
+            vendor="KUNLUNXIN",
+            priority=200,
+        ),
     ]
 
     registry.register_many(impls)