[Ascend] Integrate transformer_engine_npu && Fix the GEMM operator bug in the reference backend (#89)

## Summary

This PR adds Ascend NPU support to the TE-FL plugin system through
`torch_npu` and `transformer_engine_npu`, and fixes backward-path issues
in the reference GEMM implementation.

## Changes

### Ascend NPU backend

- Add automatic NPU availability detection and vendor-priority
registration.
- Add support for:
  - FlashAttention with SBHD, BSHD, and THD layouts
  - RMSNorm forward and backward
  - Generic and grouped GEMM
  - Multi-tensor scale and L2-norm operations
- Add THD ↔ BSHD conversion operators.
- Keep NPU dependencies lazily imported.

### Reference GEMM fixes

- Fix output shape restoration for transposed inputs.
- Do not add forward bias in backward mode.
- Compute fused bias gradients.
- Apply dGeLU using the saved forward activation.
- Preserve correct alpha scaling and 3D input behavior.

## Testing

Added coverage for:

- FlashAttention forward/backward accuracy and causal masking
- RMSNorm forward/backward
- Generic and grouped GEMM
- Multi-tensor and FP8 scale operations
- Reference GEMM backward behavior

Verified on Ascend 910C:

```text
46 passed
```

## Deps
It depends on TransformerEngineNPU. The package natively generated by
TransformerEngineNPU is named transformer_engine. Relevant packaging
logic needs to be modified so that the generated package is named
transformer_engine_npu.

Author: Darryl233

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

@@ -5,6 +5,8 @@
 from typing import Any, Optional, Tuple, Union
 import torch
 
+from .activation import dgelu_torch
+
 __all__ = [
     "general_gemm_torch",
 ]
@@ -84,25 +86,33 @@ def general_gemm_torch(
     if alpha != 1.0:
         out = out * alpha
 
-    if original_B_shape is not None:
+    # A non-transposed B contributes its outer dimensions to the output. A
+    # transposed B does not, so its flattened shape must not be restored (the
+    # latter is the layout normally used by weight-gradient GEMMs).
+    if original_B_shape is not None and not transB:
         out = out.view(original_B_shape[0], original_B_shape[1], -1)
 
     gelu_input_ret = None
-    if gelu and gelu_in is not None:
-        pass
 
-    if bias is not None:
+    # In a backward GEMM, `bias` only requests the fused BGRAD epilogue. Its
+    # value is not added to the GEMM result.
+    if bias is not None and not grad:
         if bias.device != target_device:
             bias = bias.to(target_device)
         out = out + bias
 
     if gelu:
-        if gelu_in is not None:
-            gelu_in.copy_(out)
-            gelu_input_ret = gelu_in
+        if grad:
+            if gelu_in is None:
+                raise ValueError("gelu_in must be provided for a backward GELU GEMM")
+            out = dgelu_torch(out, gelu_in, quantizer=None)
         else:
-            gelu_input_ret = out.clone()
-        out = F.gelu(out, approximate="tanh")
+            if gelu_in is not None:
+                gelu_in.copy_(out)
+                gelu_input_ret = gelu_in
+            else:
+                gelu_input_ret = out.clone()
+            out = F.gelu(out, approximate="tanh")
 
     torch_out_dtype = _convert_dtype(output_dtype)
     if torch_out_dtype is not None and out.dtype != torch_out_dtype:
@@ -121,7 +131,9 @@ def general_gemm_torch(
 
     bias_grad = None
     if grad and bias is not None:
-        pass
+        # cuBLASLt's BGRADB epilogue always reduces GEMM input B. Flattening
+        # all leading dimensions also handles sequence-shaped gradient input.
+        bias_grad = B.sum(dim=0).to(dtype=out.dtype)
 
     extra_output_ret = None
 

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

@@ -740,6 +740,52 @@ def multi_tensor_compute_scale_inv_e8m0(
             block_len,
         )
 
+    def convert_thd_to_bshd(
+        self,
+        tensor: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        b: int,
+        max_seq_len: int,
+    ) -> torch.Tensor:
+        """Convert THD (packed tokens) format to BSHD (batched, padded) format."""
+        # tensor shape: [total_tokens, num_heads, head_dim]
+        # output shape: [b, max_seq_len, num_heads, head_dim]
+        remaining_dims = tensor.shape[1:]
+        output = torch.zeros(
+            (b, max_seq_len) + remaining_dims,
+            dtype=tensor.dtype,
+            device=tensor.device,
+        )
+        for i in range(b):
+            start = cu_seqlens[i].item()
+            end = cu_seqlens[i + 1].item()
+            seq_len = end - start
+            output[i, :seq_len] = tensor[start:end]
+        return output
+
+    def convert_bshd_to_thd(
+        self,
+        tensor: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        t: int,
+    ) -> torch.Tensor:
+        """Convert BSHD (batched, padded) format to THD (packed tokens) format."""
+        # tensor shape: [b, max_seq_len, num_heads, head_dim]
+        # output shape: [t, num_heads, head_dim]
+        b = tensor.shape[0]
+        remaining_dims = tensor.shape[2:]
+        output = torch.zeros(
+            (t,) + remaining_dims,
+            dtype=tensor.dtype,
+            device=tensor.device,
+        )
+        for i in range(b):
+            start = cu_seqlens[i].item()
+            end = cu_seqlens[i + 1].item()
+            seq_len = end - start
+            output[start:end] = tensor[i, :seq_len]
+        return output
+
     def get_flash_attention_class(self):
         from .flash_attention import FlashAttentionTorch
 

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

@@ -516,6 +516,23 @@ def register_builtins(registry) -> None:
             vendor=None,
             priority=50,
         ),
+        # THD <-> BSHD format conversion
+        OpImpl(
+            op_name="convert_thd_to_bshd",
+            impl_id="reference.torch",
+            kind=BackendImplKind.REFERENCE,
+            fn=_bind_is_available(backend.convert_thd_to_bshd, is_avail),
+            vendor=None,
+            priority=50,
+        ),
+        OpImpl(
+            op_name="convert_bshd_to_thd",
+            impl_id="reference.torch",
+            kind=BackendImplKind.REFERENCE,
+            fn=_bind_is_available(backend.convert_bshd_to_thd, is_avail),
+            vendor=None,
+            priority=50,
+        ),
         # FlashAttention class getter
         OpImpl(
             op_name="get_flash_attention_class",

transformer_engine/plugin/core/backends/vendor/npu/__init__.py

@@ -0,0 +1,8 @@
+# Copyright (c) 2025, BAAI. All rights reserved.
+# Copyright (c) 2026, Huawei Technologies Co., Ltd. All rights reserved.
+#
+# See LICENSE for license information.
+
+from .npu import NPUBackend
+
+__all__ = ["NPUBackend"]

transformer_engine/plugin/core/backends/vendor/npu/flash_attention.py

@@ -0,0 +1,229 @@
+# Copyright (c) 2026, BAAI. All rights reserved.
+# Copyright (c) 2026, Huawei Technologies Co., Ltd. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""NPU Flash Attention adapter.
+
+Bridges TE-FL's FlashAttention calling convention to NPU's npu_fusion_attention kernel.
+
+TE-FL passes many parameters (qkv_layout, window_size, cp_group, fp8, etc.)
+that NPU's FlashAttention doesn't support. This adapter:
+  1. Accepts the full TE-FL parameter set
+  2. Maps qkv_layout → qkv_format (sbhd/thd)
+  3. Forwards only the supported parameters to NPU's FlashAttention
+  4. Silently ignores unsupported features (sliding window, CP, FP8, ALiBi)
+"""
+
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import torch
+
+from transformer_engine.plugin.core.ops import FlashAttentionBase
+
+
+_COMPRESSED_MASK_SIZE = 2048
+
+_COMPRESSED_CAUSAL_MASK = None
+
+
+def get_compressed_causal_mask(device="npu"):
+    global _COMPRESSED_CAUSAL_MASK
+    if _COMPRESSED_CAUSAL_MASK is None:
+        _COMPRESSED_CAUSAL_MASK = torch.triu(
+            torch.ones(
+                (_COMPRESSED_MASK_SIZE, _COMPRESSED_MASK_SIZE),
+                device=device,
+                dtype=torch.bool,
+            ),
+            diagonal=1,
+        )
+    return _COMPRESSED_CAUSAL_MASK
+
+
+class NPUFlashAttention(FlashAttentionBase):
+    """FlashAttention adapter for NPU (Ascend) hardware.
+
+    Wraps transformer_engine_npu's FlashAttention, which calls
+    torch_npu.npu_fusion_attention under the hood.
+
+    Supported features:
+      - sbhd, bshd (via transpose), and thd formats
+      - causal / padding mask types
+      - Variable-length sequences (cu_seqlens)
+      - Sparse mask optimization (via NPU's get_fa_config)
+
+    Not supported (silently ignored):
+      - Sliding window attention (window_size)
+      - ALiBi slopes
+      - Context Parallelism (cp_group, cp_stream, etc.)
+      - FP8 / quantization
+      - KV cache (inference_params)
+      - FA v2/v3 version selection
+    """
+
+    def __init__(
+        self,
+        softmax_scale: float,
+        attention_dropout: float = 0.0,
+        attention_dropout_ctx: Optional[Callable] = None,
+        attention_type: str = "self",
+        layer_number: Optional[int] = None,
+        deterministic: bool = False,
+        **kwargs,
+    ) -> None:
+        super().__init__(
+            softmax_scale=softmax_scale,
+            attention_dropout=attention_dropout,
+            attention_dropout_ctx=attention_dropout_ctx,
+            attention_type=attention_type,
+            layer_number=layer_number,
+            deterministic=deterministic,
+        )
+        self.softmax_scale = softmax_scale
+        self.attention_dropout = attention_dropout
+        self.attention_type = attention_type
+        self.layer_number = layer_number
+        self._npu_flash = None
+
+    def _ensure_backend(self):
+        """Lazy-initialize NPU FlashAttention backend."""
+        if self._npu_flash is not None:
+            return
+        from transformer_engine_npu.pytorch.attention.dot_product_attention.backends import (
+            FlashAttention as _NPUFlashAttention,
+        )
+
+        self._npu_flash = _NPUFlashAttention(self.softmax_scale)
+
+    @staticmethod
+    def _layout_to_format(qkv_layout: Optional[str]) -> str:
+        """Map TE-FL qkv_layout string to NPU qkv_format."""
+        if qkv_layout is None:
+            return "sbhd"
+        layout = qkv_layout.lower()
+        if "thd" in layout or layout.startswith("t"):
+            return "thd"
+        return "sbhd"
+
+    @staticmethod
+    def _is_bshd_layout(qkv_layout: Optional[str]) -> bool:
+        """Whether the separate Q/K/V tensors use batch-major BSHD layout."""
+        if qkv_layout is None:
+            return False
+        layout = qkv_layout.lower()
+        # Covers bs3hd, bsh3d, and bshd_bshd_bshd.
+        return layout.startswith("bs")
+
+    def _forward_impl(
+        self,
+        query_layer: torch.Tensor,
+        key_layer: torch.Tensor,
+        value_layer: torch.Tensor,
+        attention_mask: Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]] = None,
+        qkv_layout: Optional[str] = None,
+        cu_seqlens_q: Optional[torch.Tensor] = None,
+        cu_seqlens_kv: Optional[torch.Tensor] = None,
+        max_seqlen_q: Optional[int] = None,
+        max_seqlen_kv: Optional[int] = None,
+        attn_mask_type: str = "causal",
+        window_size: Optional[Tuple[int, int]] = None,
+        alibi_slopes: Optional[torch.Tensor] = None,
+        cp_group: Optional[Any] = None,
+        cp_global_ranks: Optional[List[int]] = None,
+        cp_stream: Optional[Any] = None,
+        cp_comm_type: str = "p2p",
+        fp8: bool = False,
+        fp8_meta: Optional[Dict[str, Any]] = None,
+        quantizers: Optional[Any] = None,
+        inference_params: Optional[Any] = None,
+        flash_attention_backend: Optional[Any] = None,
+        fp8_output: bool = False,
+        num_splits: Optional[int] = 1,
+    ) -> torch.Tensor:
+        """Forward pass — adapts TE-FL args to NPU FlashAttention interface.
+
+        Only passes: query, key, value, attention_mask, qkv_format,
+                     cu_seqlens_q, cu_seqlens_kv, attn_mask_type
+
+        Raises:
+            NotImplementedError: For features that would produce incorrect results
+                if silently ignored (window_size, alibi_slopes, cp_group).
+
+        Warns:
+            For features that don't affect correctness but differ from user
+            expectation (fp8, inference_params).
+        """
+        # --- Validate: features that would silently produce wrong results ---
+        if window_size is not None and window_size not in ((-1, -1), (-1, 0)):
+            raise NotImplementedError(
+                "NPU FlashAttention does not support sliding window attention "
+                f"(window_size={window_size}). npu_fusion_attention only computes "
+                "full causal/padding attention. Either disable sliding window or "
+                "use UnfusedDotProductAttention as fallback."
+            )
+
+        if alibi_slopes is not None:
+            raise NotImplementedError(
+                "NPU FlashAttention does not support ALiBi position encoding "
+                "(alibi_slopes). npu_fusion_attention has no ALiBi parameter. "
+                "Use RoPE or other position encoding supported by NPU."
+            )
+
+        if cp_group is not None:
+            raise NotImplementedError(
+                "NPU FlashAttention does not support Context Parallelism "
+                "(cp_group). Ring attention / CP requires NPU-specific HCCL "
+                "implementation which is not yet available."
+            )
+
+        # --- Warn: features that don't break correctness but differ from expectation ---
+        if fp8:
+            import warnings
+
+            warnings.warn(
+                "NPU FlashAttention does not support FP8 attention computation. "
+                "Falling back to BF16/FP16 precision. Results are correct but "
+                "without FP8 performance optimization.",
+                stacklevel=2,
+            )
+
+        if inference_params is not None:
+            import warnings
+
+            warnings.warn(
+                "NPU FlashAttention does not support KV cache (inference_params). "
+                "Full recomputation will be used. This is correct but slower for "
+                "autoregressive inference.",
+                stacklevel=2,
+            )
+
+        # TransformerEngineNPU only accepts sequence-major SBHD or packed THD.
+        # Convert batch-major BSHD inputs explicitly instead of only relabeling
+        # their layout, which would swap the semantic batch and sequence axes.
+        input_is_bshd = self._is_bshd_layout(qkv_layout)
+        if input_is_bshd:
+            query_layer = query_layer.transpose(0, 1).contiguous()
+            key_layer = key_layer.transpose(0, 1).contiguous()
+            value_layer = value_layer.transpose(0, 1).contiguous()
+
+        qkv_format = self._layout_to_format(qkv_layout)
+        if attn_mask_type in ("causal", "padding_causal", "padding,causal", "causal,padding"):
+            attention_mask = get_compressed_causal_mask(query_layer.device)
+
+        self._ensure_backend()
+        output = self._npu_flash(
+            query_layer,
+            key_layer,
+            value_layer,
+            attention_mask=attention_mask,
+            qkv_format=qkv_format,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_kv=cu_seqlens_kv,
+            attn_mask_type=attn_mask_type,
+        )
+
+        if input_is_bshd:
+            output = output.transpose(0, 1).contiguous()
+
+        return output