Merge branch 'flagos-ai:main' into main

Author: BrianPei

transformer_engine/__init__.py

@@ -10,6 +10,35 @@
 from importlib import metadata
 import transformer_engine.common
 
+import torch
+
+# Public, simple global (kept for backward compatibility).
+TE_DEVICE_TYPE = "cuda"
+TE_PLATFORM = torch.cuda
+
+# Apply MUSA (VENDOR) Patches, such as torch.cuda.device -> torch.musa.device
+try:
+    from .plugin.core.backends.vendor.musa.patches import apply_patch as _musa_apply_patch
+
+    _musa_apply_patch()
+except Exception as e:
+    pass
+
+
+def te_device_type(default: str = "cuda") -> str:
+    try:
+        return TE_DEVICE_TYPE
+    except Exception:
+        return default
+
+
+def te_platform(default=torch.cuda):
+    try:
+        return TE_PLATFORM
+    except Exception:
+        return default
+
+
 try:
     from . import pytorch
 except ImportError:

transformer_engine/debug/features/fake_quant.py

@@ -14,6 +14,7 @@
 
 
 import transformer_engine_torch as tex
+from transformer_engine import te_device_type
 from transformer_engine.debug.features.api import TEConfigAPIMapper
 from transformer_engine.common.recipe import Format
 from transformer_engine.pytorch.tensor import Quantizer
@@ -30,7 +31,9 @@ def fake_quantize(tensor: torch.Tensor, fp8_format: tex.DType, out=None):
         torch.float16,
         torch.bfloat16,
     ), "[NVTORCH INSPECT ERROR] Unsupported tensor type."
-    assert tensor.is_cuda, "[NVTORCH INSPECT ERROR] Must be a GPU tensor."
+    assert (
+        tensor.device.type == te_device_type()
+    ), f"[NVTORCH INSPECT ERROR] Must be a {te_device_type()} tensor."
     assert fp8_format in {
         "FP8E4M3",
         "FP8E5M2",

transformer_engine/debug/features/log_fp8_tensor_stats.py

@@ -14,6 +14,7 @@
 from nvdlfw_inspect.debug_features.log_tensor_stats import LogTensorStats as BaseLogTensorStats
 from nvdlfw_inspect.registry import Registry, api_method
 
+from transformer_engine import te_device_type
 from transformer_engine.debug.features.utils.stats_buffer import STATS_BUFFERS
 from transformer_engine.pytorch.tensor import Quantizer, QuantizedTensor
 from transformer_engine.pytorch.tensor.float8_tensor import (
@@ -47,7 +48,10 @@ def _get_new_quantizer(recipe_name, fp8_dtype):
         return Float8BlockQuantizer(fp8_dtype=fp8_dtype, rowwise=True, columnwise=True)
     if recipe_name == "fp8_current_scaling":
         return Float8CurrentScalingQuantizer(
-            fp8_dtype=fp8_dtype, device=torch.device("cuda"), rowwise=True, columnwise=True
+            fp8_dtype=fp8_dtype,
+            device=torch.device(te_device_type()),
+            rowwise=True,
+            columnwise=True,
         )
     if recipe_name == "mxfp8":
         return MXFP8Quantizer(fp8_dtype=fp8_dtype, rowwise=True, columnwise=True)

transformer_engine/debug/features/per_tensor_scaling.py

@@ -11,7 +11,9 @@
 import nvdlfw_inspect.api as debug_api
 from nvdlfw_inspect.registry import Registry, api_method
 
+
 import transformer_engine_torch as tex
+from transformer_engine import te_device_type
 from transformer_engine.pytorch.tensor import Quantizer
 from transformer_engine.pytorch.tensor.float8_tensor import (
     Float8Tensor,
@@ -33,7 +35,9 @@ def per_tensor_cast(
         torch.float16,
         torch.bfloat16,
     ), "[NVTORCH INSPECT ERROR] Unsupported tensor type for per tensor current scaling"
-    assert tensor.is_cuda, "[NVTORCH INSPECT ERROR] Must be a GPU tensor."
+    assert (
+        tensor.device.type == te_device_type()
+    ), f"[NVTORCH INSPECT ERROR] Must be a {te_device_type()} tensor."
     assert fp8_dtype in {
         tex.DType.kFloat8E4M3,
         tex.DType.kFloat8E5M2,

transformer_engine/debug/features/utils/stats_buffer.py

@@ -16,6 +16,7 @@
 from nvdlfw_inspect.utils import gather_along_first_dim
 from nvdlfw_inspect.logging import MetricLogger
 
+from transformer_engine import te_device_type
 from transformer_engine.debug.features.utils.stats_computation import (
     STATS,
     DEPENDENCIES,
@@ -41,14 +42,14 @@ def __init__(self, layer_name, tensor_name, stats, reduction_group, reduce_withi
         for stat in stats:
             self.stats_to_compute = self.stats_to_compute | DEPENDENCIES[stat]
 
-        self._buffer = torch.zeros(len(STATS), dtype=torch.float32).cuda()
+        self._buffer = torch.zeros(len(STATS), dtype=torch.float32).to(te_device_type())
         self._new_buffer = self._buffer.clone()
         self._tmp_buffer = self._buffer.clone()
 
         # in case of data parallelism it is possible that layer will not be run on one node
         # modified is set to True if node is run
         # we do not take not run nodes into account
-        self.modified = torch.tensor([False], dtype=torch.bool).cuda()
+        self.modified = torch.tensor([False], dtype=torch.bool).to(te_device_type())
         self.iteration = None
         self.skip_reduction = False
 

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

@@ -9,6 +9,7 @@
 from packaging.version import Version as PkgVersion
 
 import torch
+from transformer_engine import te_device_type
 from transformer_engine.pytorch.utils import (
     get_device_compute_capability,
 )
@@ -283,8 +284,10 @@ def _forward_impl(
             for x in [query_layer, key_layer, value_layer]
         ), "FLAttention only supports FP16 and BF16 data types, or Float8Tensors."
         assert (
-            query_layer.is_cuda and key_layer.is_cuda and value_layer.is_cuda
-        ), "FLAttention only supports CUDA tensors."
+            query_layer.device.type == te_device_type()
+            and key_layer.device.type == te_device_type()
+            and value_layer.device.type == te_device_type()
+        ), f"FLAttention only supports {te_device_type()} tensors."
         assert qkv_layout in QKVLayouts, f"FLAttention does not support qkv_layout = {qkv_layout}!"
 
         cp_size = 1

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

@@ -17,6 +17,9 @@
     multi_tensor_adam_param_remainder_fl,
     multi_tensor_l2_norm_fl,
     generic_gemm_fl,
+    scaled_masked_softmax_forward_fl,
+    scaled_masked_softmax_backward_fl,
+    te_general_grouped_gemm_fl,
 )
 
 
@@ -116,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,
@@ -160,6 +203,23 @@ def rmsnorm_bwd(
     def get_fused_attn_backend(self, *args, **kwargs) -> int:
         return NVTE_Fused_Attn_Backend.NVTE_No_Backend
 
+    # Softmax functions
+    def scaled_masked_softmax_forward(
+        self,
+        input: torch.Tensor,
+        mask: torch.Tensor,
+        scale_factor: Union[float, torch.Tensor],
+    ) -> torch.Tensor:
+        return scaled_masked_softmax_forward_fl(input, mask, scale_factor)
+
+    def scaled_masked_softmax_backward(
+        self,
+        output_grad_: torch.Tensor,
+        softmax_results_: torch.Tensor,
+        scale_factor: float,
+    ) -> torch.Tensor:
+        return scaled_masked_softmax_backward_fl(output_grad_, softmax_results_, scale_factor)
+
     # multi-tensor functions
     def multi_tensor_scale(
         self,
@@ -243,7 +303,7 @@ def get_cudnn_version(self) -> int:
         return 90000
 
     def get_num_cublas_streams(self) -> int:
-        return 0
+        return 4  # keep consistent with transformer_engine/common/util/multi_stream.cpp, get_num_compute_streams()
 
     ############## class func #################################
     def get_flash_attention_class(self):

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

@@ -6,3 +6,4 @@
 from .rmsnorm import *
 from .fused_adam import *
 from .multi_tensor import *
+from .softmax import *

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