Add python-level patches to supporting multiple platforms

transformer_engine/__init__.py

@@ -10,6 +10,37 @@
 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()
+    print("[TE-FL] MUSA patches applied")
+except Exception as e:
+    print(f"[TE-FL] MUSA patches not applied: {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

@@ -243,7 +243,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/register_ops.py

@@ -124,6 +124,22 @@ def register_builtins(registry) -> None:
             vendor=None,
             priority=150,
         ),
+        OpImpl(
+            op_name="get_num_cublas_streams",
+            impl_id="default.flagos",
+            kind=BackendImplKind.DEFAULT,
+            fn=_bind_is_available(backend.get_num_cublas_streams, is_avail),
+            vendor=None,
+            priority=150,
+        ),
+        OpImpl(
+            op_name="get_cudnn_version",
+            impl_id="default.flagos",
+            kind=BackendImplKind.DEFAULT,
+            fn=_bind_is_available(backend.get_cudnn_version, is_avail),
+            vendor=None,
+            priority=150,
+        ),
     ]
 
     registry.register_many(impls)

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

@@ -44,20 +44,35 @@ def scaled_masked_softmax_forward_torch(
     mask: torch.Tensor,
     scale: float,
 ) -> torch.Tensor:
-    # Handle uint8 mask (CUDA format: 1=masked, 0=unmasked)
-    # Convert to additive mask (-10000 for masked positions, 0 for unmasked)
-    if mask.dtype == torch.uint8:
-        additive_mask = torch.zeros_like(input, dtype=input.dtype)
-        # Expand mask if needed (mask shape: batch, 1, seq_q, seq_k)
-        if mask.dim() == 4 and mask.size(1) == 1 and input.dim() == 4:
-            mask = mask.expand_as(input)
-        additive_mask = additive_mask.masked_fill(mask.bool(), -10000.0)
-    else:
-        additive_mask = mask
-
-    scaled_input = input * scale + additive_mask
-
-    return F.softmax(scaled_input, dim=-1)
+    """Reference forward matching TE CUDA `scaled_masked_softmax_warp_forward`.
+
+    Integer/bool mask (same as uint8 kernel contract):
+      - **Exactly** ``mask == 1`` means **masked** (logit set to ``-10000``, not ``input*scale`` offset).
+      - Any other value (typically 0) means **unmasked** (logit is ``input * scale``).
+
+    Floating mask: treated as **additive** bias in logit space (already scaled), added after
+    ``input * scale``.
+
+    Common pitfalls this avoids vs the old implementation:
+    1) ``input * scale + (-10000)`` on masked positions ≠ CUDA's plain ``-10000``.
+    2) Non-uint8 masks (bool, int) were used as direct addends → wrong (0/1 added to logits).
+    3) ``mask.bool()`` masks any nonzero byte; CUDA only masks when ``mask == 1``.
+    """
+    if mask.dim() == 4 and mask.size(1) == 1 and input.dim() == 4:
+        mask = mask.expand_as(input)
+
+    scaled = input * scale
+
+    if mask.is_floating_point():
+        scaled = scaled + mask.to(dtype=scaled.dtype)
+        return F.softmax(scaled, dim=-1)
+
+    # Integer / bool: align with CUDA (masked iff value == 1)
+    scaled = scaled.masked_fill(mask == 1, -10000.0)
+    # CUDA zeros output row when every position in the softmax dim is masked (max == -10000)
+    all_masked = (mask == 1).all(dim=-1, keepdim=True)
+    out = F.softmax(scaled, dim=-1)
+    return out.masked_fill(all_masked, 0.0)
 
 
 def scaled_masked_softmax_backward_torch(

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

@@ -476,7 +476,7 @@ def get_cudnn_version(self) -> int:
         return 0
 
     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()
 
     # Multi-tensor functions
     def multi_tensor_scale(

transformer_engine/plugin/core/backends/vendor/musa/patches.py

@@ -0,0 +1,72 @@
+"""Python-side compatibility patches for the MUSA vendor backend."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+import torch
+
+
+def _noop(*args, **kwargs):
+    return None
+
+
+# Patches: (parent_object, attribute_name, replacement_callable)
+_PATCH_CALLS: list[tuple[object, str, Callable[..., object]]] = [
+    # We do not recommend replace is_available, due to its device-related behavior.
+    # (torch.cuda, "is_available", torch.musa.is_available),
+    (torch.cuda, "get_device_properties", torch.musa.get_device_properties),
+    (torch.cuda, "device", torch.musa.device),
+    (torch.cuda, "current_device", torch.musa.current_device),
+    (torch.cuda, "synchronize", torch.musa.synchronize),
+    (torch.cuda, "is_current_stream_capturing", torch.musa.is_current_stream_capturing),
+    # TODO: Add NVTX patches for MUSA.
+    # NVTX is CUDA-specific; make it a no-op on MUSA.
+    (torch.cuda.nvtx, "range_push", _noop),
+    (torch.cuda.nvtx, "range_pop", _noop),
+    # TODO: Add other patches for MUSA.
+]
+
+
+def apply_patch() -> None:
+    """Apply MUSA Python-side patches (idempotent, best-effort)."""
+    try:
+        from .musa import MUSABackend
+
+        if not MUSABackend().is_available():
+            return
+    except Exception as e:
+        print(f"[TE-FL] MUSA backend not available: {e}")
+        # If backend availability can't be determined, don't patch.
+        return
+
+    # Mark TE global device type for Python-side callers.
+    # IMPORTANT: do not import `transformer_engine` here, because TE's `__init__.py`
+    # imports this module to run patches and that would cause a circular import.
+    try:
+        import transformer_engine
+
+        transformer_engine.TE_DEVICE_TYPE = "musa"
+        transformer_engine.TE_PLATFORM = torch.musa
+    except Exception as e:
+        print(f"[TE-FL Musa Patches] Error setting TE device type or platform: {e}")
+        # Best-effort: don't fail patching if we can't set the global.
+        pass
+
+    # Only patch when torch.musa exists and is usable.
+    if not hasattr(torch, "musa"):
+        return
+    try:
+        if not torch.musa.is_available():
+            return
+    except Exception:
+        return
+
+    for parent, attr, replacement in _PATCH_CALLS:
+        if not hasattr(parent, attr):
+            continue
+        try:
+            setattr(parent, attr, replacement)
+        except Exception:
+            # Best-effort: patching should never crash import/initialization.
+            continue

transformer_engine/pytorch/attention/dot_product_attention/backends.py

@@ -15,6 +15,7 @@
 import torch
 import torch.nn.functional as F
 import transformer_engine_torch as tex
+from transformer_engine import te_device_type
 from transformer_engine.pytorch.utils import (
     get_device_compute_capability,
     split_tensor_along_dim,
@@ -387,10 +388,10 @@ def forward(
                 fp8_recipe = fp8_meta["local_recipes"][0]
             if fp8_recipe.float8_current_scaling():
                 S_quantizer = Float8CurrentScalingQuantizer(
-                    fp8_dtype=S_quantizer.dtype, device="cuda"
+                    fp8_dtype=S_quantizer.dtype, device=te_device_type()
                 )
                 dP_quantizer = Float8CurrentScalingQuantizer(
-                    fp8_dtype=dP_quantizer.dtype, device="cuda"
+                    fp8_dtype=dP_quantizer.dtype, device=te_device_type()
                 )
 
             if "2" in qkv_layout or "3" in qkv_layout:
@@ -676,8 +677,10 @@ def forward(
             for x in [query_layer, key_layer, value_layer]
         ), "FlashAttention only supports FP16 and BF16 data types, or Float8Tensors."
         assert (
-            query_layer.is_cuda and key_layer.is_cuda and value_layer.is_cuda
-        ), "FlashAttention currently 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"FlashAttention currently only supports {te_device_type()} tensors."
         assert (
             qkv_layout in QKVLayouts
         ), f"FlashAttention does not support qkv_layout = {qkv_layout}!"
@@ -1738,8 +1741,10 @@ def forward(
             for x in [query_layer, key_layer, value_layer]
         ), "FusedAttention only supports FP16 and BF16 data types, or Float8Tensors."
         assert (
-            query_layer.is_cuda and key_layer.is_cuda and value_layer.is_cuda
-        ), "FusedAttention 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"FusedAttention only supports {te_device_type()} tensors."
         assert (
             qkv_layout in QKVLayouts
         ), f"FusedAttention does not support qkv_layout = {qkv_layout}!"

transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py

@@ -14,6 +14,7 @@
 from torch.nn.parameter import Parameter
 
 import transformer_engine_torch as tex
+from transformer_engine import te_device_type
 from transformer_engine.common.recipe import (
     Format,
     Recipe,
@@ -420,12 +421,14 @@ def __init__(
             self.softmax_offset = None
         if self.softmax_type == "off-by-one":
             self.softmax_offset = torch.zeros(
-                self.num_attention_heads // self.tp_size, device="cuda"
+                self.num_attention_heads // self.tp_size, device=te_device_type()
             )
         if self.softmax_type == "learnable":
             self.register_parameter(
                 "softmax_offset",
-                Parameter(torch.empty(self.num_attention_heads // self.tp_size, device="cuda")),
+                Parameter(
+                    torch.empty(self.num_attention_heads // self.tp_size, device=te_device_type())
+                ),
                 get_rng_state_tracker=get_rng_state_tracker,
             )
 
@@ -1026,8 +1029,10 @@ def forward(
 
             # checks for q/k/v shapes
             assert (
-                query_layer.is_cuda and key_layer.is_cuda and value_layer.is_cuda
-            ), "DotProductAttention 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"DotProductAttention only supports {te_device_type()} tensors."
             assert (
                 query_layer.dtype == key_layer.dtype and query_layer.dtype == value_layer.dtype
             ), "Queries, keys and values must have the same data type!"