Abstract CUDA hardcodes into configurable te_device_type / te_platform

transformer_engine/__init__.py

@@ -10,6 +10,26 @@
 from importlib import metadata
 import transformer_engine.common
 
+
+# Plugin system: set NVTE_PLUGIN to the plugin module name to enable.
+# e.g. NVTE_PLUGIN=transformer_engine_plugin_fl
+_nvte_plugin = os.environ.get("NVTE_PLUGIN")
+if _nvte_plugin:
+    try:
+        from importlib import import_module
+
+        _patches = import_module(f"{_nvte_plugin}.patches")
+        _patches.apply_patches()
+    except Exception as e:
+        import warnings
+
+        warnings.warn(
+            f"NVTE_PLUGIN={_nvte_plugin} but plugin patch apply failed: {e}",
+            RuntimeWarning,
+            stacklevel=1,
+        )
+
+
 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

@@ -15,6 +15,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.debug.features.utils import get_reduction_params, next_enabled_iter
 from transformer_engine.pytorch.tensor import Quantizer, QuantizedTensor
@@ -58,7 +59,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

@@ -12,6 +12,7 @@
 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 +34,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/pytorch/__init__.py

@@ -9,13 +9,48 @@
 import functools
 
 import torch
+import transformer_engine
 
 from transformer_engine.common import load_framework_extension
 from transformer_engine.pytorch.torch_version import torch_version
 
 assert torch_version() >= (2, 1), f"Minimum torch version 2.1 required. Found {torch_version()}."
 
 load_framework_extension("torch")
+
+# Device type and platform globals — live here because torch cannot be imported
+# at the top-level transformer_engine package (shared with JAX).
+# Only set defaults if the plugin hasn't already patched them.
+if not hasattr(transformer_engine, "TE_DEVICE_TYPE"):
+    transformer_engine.TE_DEVICE_TYPE = "cuda"
+TE_DEVICE_TYPE = transformer_engine.TE_DEVICE_TYPE
+
+
+def te_device_type(default: str = "cuda") -> str:
+    try:
+        return transformer_engine.TE_DEVICE_TYPE
+    except Exception:
+        return default
+
+
+if not hasattr(transformer_engine, "te_device_type"):
+    transformer_engine.te_device_type = te_device_type
+
+if not hasattr(transformer_engine, "TE_PLATFORM"):
+    transformer_engine.TE_PLATFORM = torch.cuda
+TE_PLATFORM = transformer_engine.TE_PLATFORM
+
+
+def te_platform(default=torch.cuda):
+    try:
+        return transformer_engine.TE_PLATFORM
+    except Exception:
+        return default
+
+
+if not hasattr(transformer_engine, "te_platform"):
+    transformer_engine.te_platform = te_platform
+
 from transformer_engine.pytorch.module import LayerNormLinear
 from transformer_engine.pytorch.module import Linear
 from transformer_engine.pytorch.module import LayerNormMLP

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,
@@ -456,10 +457,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:
@@ -750,8 +751,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}!"
@@ -1821,8 +1824,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,
@@ -434,12 +435,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.zeros(self.num_attention_heads // self.tp_size, device="cuda")),
+                Parameter(
+                    torch.zeros(self.num_attention_heads // self.tp_size, device=te_device_type())
+                ),
                 get_rng_state_tracker=get_rng_state_tracker,
             )
 
@@ -1057,8 +1060,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!"

transformer_engine/pytorch/attention/dot_product_attention/softmax.py

@@ -8,6 +8,7 @@
 import torch
 from torch import nn
 import transformer_engine_torch as tex
+from transformer_engine import te_device_type
 from transformer_engine.pytorch.export import is_in_onnx_export_mode
 
 
@@ -24,7 +25,7 @@ def _get_default_causal_mask(mask_type: str, sq: int, sk: int) -> torch.Tensor:
     def _get_mask():
         diagonal_offset = sk - sq + 1 if "bottom_right" in mask_type else 1
         return torch.triu(
-            torch.ones(sq, sk, dtype=torch.bool, device="cuda"), diagonal=diagonal_offset
+            torch.ones(sq, sk, dtype=torch.bool, device=te_device_type()), diagonal=diagonal_offset
         )
 
     if is_in_onnx_export_mode():

transformer_engine/pytorch/attention/dot_product_attention/utils.py

@@ -21,6 +21,7 @@
 import torch.nn.functional as F
 import transformer_engine_torch as tex
 import transformer_engine as te
+from transformer_engine import te_device_type
 from transformer_engine.pytorch.cpp_extensions.fused_attn import (
     QKVLayout,
     AttnBiasType,
@@ -1103,6 +1104,10 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt
             logger.debug("Disabling FusedAttention for determinism reasons with post_scale_bias")
             use_fused_attention = False
             fused_attention_backend = None
+        if is_training and device_compute_capability >= (10, 0):
+            logger.debug("Disabling FusedAttention for determinism reasons on Blackwell")
+            use_fused_attention = False
+            fused_attention_backend = None
 
     # use_flash_attention may have been set above
     use_flash_attention_2 = use_flash_attention and use_flash_attention_2
@@ -1243,13 +1248,13 @@ def get_padding_mask(
                 ],
                 dim=0,
             )
-    attention_mask_q = attention_mask_q.to(device="cuda")
+    attention_mask_q = attention_mask_q.to(device=te_device_type())
     if attention_type == "self":
         attention_mask = attention_mask_q
     else:
         attention_mask = (
             attention_mask_q,
-            attention_mask_kv.to(device="cuda"),
+            attention_mask_kv.to(device=te_device_type()),
         )
     return attention_mask
 
@@ -1368,9 +1373,11 @@ def get_full_mask(
         actual_seqlens_kv = m[:, 0, 0, :].sum(dim=1)
 
     # apply SWA mask
-    mask = torch.arange(max_seqlen_q, dtype=torch.int32, device="cuda").view(
+    mask = torch.arange(max_seqlen_q, dtype=torch.int32, device=te_device_type()).view(
         1, 1, max_seqlen_q, 1
-    ) - torch.arange(max_seqlen_kv, dtype=torch.int32, device="cuda").view(1, 1, 1, max_seqlen_kv)
+    ) - torch.arange(max_seqlen_kv, dtype=torch.int32, device=te_device_type()).view(
+        1, 1, 1, max_seqlen_kv
+    )
     swa_left = None
     swa_right = None
     if attn_mask_type == "causal_bottom_right" or (
@@ -1466,7 +1473,7 @@ def get_alibi(
                 m_hat = torch.pow(m_hat_0, torch.arange(1, 1 + 2 * (num_heads - n), 2))
                 m = torch.cat([m, m_hat])
 
-            _alibi_cache["_alibi_slopes"] = m.to(dtype=torch.float32, device="cuda")
+            _alibi_cache["_alibi_slopes"] = m.to(dtype=torch.float32, device=te_device_type())
         _alibi_cache["_num_heads"] = num_heads
         _alibi_cache["_alibi_slopes_require_update"] = False
 
@@ -1479,9 +1486,9 @@ def get_alibi(
         else:
             raise ValueError("ALiBi slopes cannot exceed 2 dimensions.")
 
-        bias = torch.arange(max_seqlen_q, dtype=torch.int32, device="cuda").view(
+        bias = torch.arange(max_seqlen_q, dtype=torch.int32, device=te_device_type()).view(
             1, 1, max_seqlen_q, 1
-        ) - torch.arange(max_seqlen_kv, dtype=torch.int32, device="cuda").view(
+        ) - torch.arange(max_seqlen_kv, dtype=torch.int32, device=te_device_type()).view(
             1, 1, 1, max_seqlen_kv
         )
         if actual_seqlens_q is None and actual_seqlens_kv is None:
@@ -1501,7 +1508,9 @@ def get_alibi(
         _alibi_cache["_max_seqlen_q"], _alibi_cache["_max_seqlen_kv"] = max_seqlen_q, max_seqlen_kv
         _alibi_cache["_bottom_right_alignment"] = bottom_right_alignment
         bias_dtype = torch.float32 if bias_dtype is None else bias_dtype
-        _alibi_cache["_alibi_bias"] = bias.contiguous().to(dtype=bias_dtype, device="cuda")
+        _alibi_cache["_alibi_bias"] = bias.contiguous().to(
+            dtype=bias_dtype, device=te_device_type()
+        )
         _alibi_cache["_alibi_bias_require_update"] = False
 
     return _alibi_cache["_alibi_slopes"], _alibi_cache["_alibi_bias"]
@@ -1516,7 +1525,7 @@ def get_cu_seqlens(mask: torch.Tensor) -> torch.Tensor:
     mask = mask.squeeze(1).squeeze(1)
     reduced_mask = mask.logical_not().sum(dim=1)
     cu_seqlens = reduced_mask.cumsum(dim=0).to(torch.int32)
-    zero = torch.zeros(1, dtype=torch.int32, device="cuda")
+    zero = torch.zeros(1, dtype=torch.int32, device=te_device_type())
     cu_seqlens = torch.cat((zero, cu_seqlens))
 
     return cu_seqlens
@@ -1534,7 +1543,7 @@ def get_cu_seqlens_and_indices(mask: torch.Tensor) -> Tuple[torch.Tensor, torch.
 
     reduced_mask = mask.logical_not().sum(dim=1)
     cu_seqlens = reduced_mask.cumsum(dim=0).to(torch.int32)
-    zero = torch.zeros(1, dtype=torch.int32, device="cuda")
+    zero = torch.zeros(1, dtype=torch.int32, device=te_device_type())
     cu_seqlens = torch.cat((zero, cu_seqlens))
 
     mask = mask.reshape(-1)
@@ -1559,7 +1568,12 @@ def get_indices(max_seqlen: int, cu_seqlens: torch.Tensor) -> torch.Tensor:
     bs = len(cu_seqlens) - 1
     seqlens = cu_seqlens[1:] - cu_seqlens[:-1]
     indices = [i * max_seqlen + ii for i, j in enumerate(seqlens) for ii in range(j)]
-    indices = torch.Tensor(indices).unsqueeze(1).unsqueeze(1).to(dtype=torch.int64, device="cuda")
+    indices = (
+        torch.Tensor(indices)
+        .unsqueeze(1)
+        .unsqueeze(1)
+        .to(dtype=torch.int64, device=te_device_type())
+    )
 
     num_nonzeros = indices.shape[0]
     pad_amount = bs * max_seqlen - num_nonzeros

transformer_engine/pytorch/attention/inference.py

@@ -11,6 +11,7 @@
 import torch
 
 import transformer_engine_torch as tex
+from transformer_engine import te_device_type
 from transformer_engine.pytorch.cpp_extensions.fused_attn import QKVFormat
 
 __all__ = ["InferenceParams", "KVCacheManager", "NonPagedKVCacheManager", "PagedKVCacheManager"]
@@ -626,7 +627,7 @@ def __init__(
         self.allocated_pages = defaultdict(list)
         # page table, [batch_size, max_pages_per_seq]
         self.page_table = torch.zeros(
-            self.max_batch_size, self.max_pages_per_seq, dtype=torch.int32, device="cuda"
+            self.max_batch_size, self.max_pages_per_seq, dtype=torch.int32, device=te_device_type()
         )
 
     def reset(self):