Abstract CUDA hardcodes into configurable te_device_type / te_platform

transformer_engine/common/__init__.py

@@ -231,6 +231,14 @@ def load_framework_extension(framework: str) -> None:
             # plugin module may import transformer_engine_torch_nv at top level.
             sys.modules[module_name + "_nv"] = solib
             _plugin = importlib.import_module(_nvte_plugin)
+            # Apply device patches before Transformer Engine's PyTorch submodules
+            # capture references to torch.cuda entry points. The submodule is
+            # optional: a plugin that needs no patches may omit it, and a
+            # single-file plugin has no submodules to search.
+            if hasattr(_plugin, "__path__") and (
+                importlib.util.find_spec(f"{_nvte_plugin}.patches") is not None
+            ):
+                importlib.import_module(f"{_nvte_plugin}.patches").apply_patches()
             _plugin.load_plugins()
         except Exception as e:  # pylint: disable=broad-exception-caught
             # Rollback to pre-plugin state if plugin failed to fully initialize

transformer_engine/debug/features/fake_quant.py

@@ -13,6 +13,7 @@
 from nvdlfw_inspect.utils import append_parent_docstring
 
 
+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 import DType
@@ -30,7 +31,9 @@ def fake_quantize(tensor: torch.Tensor, fp8_format: str, 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

@@ -17,6 +17,8 @@
 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
+
+from transformer_engine import te_device_type
 from transformer_engine.pytorch.tensor.float8_tensor import (
     Float8Quantizer,
     Float8CurrentScalingQuantizer,
@@ -57,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

@@ -11,6 +11,7 @@
 import nvdlfw_inspect.api as debug_api
 from nvdlfw_inspect.registry import Registry, api_method
 
+from transformer_engine import te_device_type
 from transformer_engine.pytorch import DType
 from transformer_engine.pytorch.tensor import Quantizer
 from transformer_engine.pytorch.tensor.float8_tensor import (
@@ -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 {
         DType.kFloat8E4M3,
         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,34 @@
 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).
+# Use plugin-set values if available, otherwise fall back to CUDA defaults.
+transformer_engine.TE_DEVICE_TYPE = getattr(transformer_engine, "TE_DEVICE_TYPE", "cuda")
+transformer_engine.TE_PLATFORM = getattr(transformer_engine, "TE_PLATFORM", torch.cuda)
+
+
+def te_device_type() -> str:
+    """Torch device type used by Transformer Engine, e.g. `cuda`."""
+    return transformer_engine.TE_DEVICE_TYPE
+
+
+def te_platform():
+    """Torch device module used by Transformer Engine, e.g. `torch.cuda`."""
+    return transformer_engine.TE_PLATFORM
+
+
+transformer_engine.te_device_type = te_device_type
+transformer_engine.te_platform = te_platform
 from transformer_engine.pytorch import constants
 from transformer_engine.pytorch.constants import CPLoadBalancingStrategy, DType
 from transformer_engine.pytorch.module import LayerNormLinear

transformer_engine/pytorch/attention/dot_product_attention/backends.py

@@ -18,6 +18,8 @@
 
 import torch
 import torch.nn.functional as F
+
+from transformer_engine import te_device_type
 from transformer_engine.pytorch.utils import (
     get_device_compute_capability,
     split_tensor_along_dim,
@@ -663,10 +665,10 @@ def _forward(
             # S/dP are forced to use DS quantizers in DPA.init_fp8_metadata; revert them here for true CS emulation
             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()
                 )
             # disable swizzle for MXFP8Quantizer
             for quantizer in [

transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py

@@ -8,6 +8,7 @@
 import torch
 import transformer_engine_torch as tex
 
+from transformer_engine import te_device_type
 from transformer_engine.pytorch.utils import (
     get_cudnn_version,
     nvtx_range_pop,
@@ -299,7 +300,9 @@ def get_thd_partitioned_indices(
             cp_size,
         )
         target_device = torch.device(device if device is not None else cu_seqlens_padded.device)
-        target_dtype = torch.int32 if target_device.type == "cuda" else cu_seqlens_padded.dtype
+        target_dtype = (
+            torch.int32 if target_device.type == te_device_type() else cu_seqlens_padded.dtype
+        )
         chunk_size = total_tokens // cp_size
         return torch.arange(
             cp_rank * chunk_size,

transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py

@@ -15,6 +15,7 @@
 from torch.fx.experimental.symbolic_shapes import guard_scalar
 from torch.nn.parameter import Parameter
 
+from transformer_engine import te_device_type
 from transformer_engine.common.recipe import (
     Format,
     Recipe,
@@ -962,12 +963,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,
             )
 

transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py

@@ -10,6 +10,7 @@
 from typing import Any, Callable, Dict, Optional, Tuple
 
 import torch
+from transformer_engine import te_device_type
 
 _cudnn_score_mod_handles: Dict[torch.device, Any] = {}
 _cudnn_score_mod_graph_cache: Dict[Tuple[Any, ...], Any] = {}
@@ -140,7 +141,7 @@ def _score_mod_callback_cache_key(callback: Optional[Callable]) -> Any:
 
 def _score_mod_device_key(device: torch.device) -> Tuple[Any, ...]:
     """Normalize a tensor device for graph cache keys."""
-    if device.type == "cuda":
+    if device.type == te_device_type():
         index = device.index
         if index is None:
             index = torch.cuda.current_device()
@@ -194,10 +195,10 @@ def _wrapped_score_mod(sdpa_graph, score_tensor):
 
 def _get_cudnn_current_stream_handle(cudnn, device: torch.device):
     """Return a cuDNN handle for device, bound to PyTorch's current stream."""
-    if device.type != "cuda":
+    if device.type != te_device_type():
         raise ValueError(f"Flex Attention only supports CUDA tensors, got device {device}.")
     if device.index is None:
-        device = torch.device("cuda", torch.cuda.current_device())
+        device = torch.device(te_device_type(), torch.cuda.current_device())
 
     handle = _cudnn_score_mod_handles.get(device)
     with torch.cuda.device(device):
@@ -287,8 +288,8 @@ def _execute_cudnn_graph(
     """Execute a built cuDNN frontend Python graph."""
     cudnn = _import_cudnn_frontend()
 
-    if device.type == "cuda" and device.index is None:
-        device = torch.device("cuda", torch.cuda.current_device())
+    if device.type == te_device_type() and device.index is None:
+        device = torch.device(te_device_type(), torch.cuda.current_device())
     workspace = torch.empty(
         workspace_size,
         device=device,

transformer_engine/pytorch/attention/dot_product_attention/softmax.py

@@ -10,6 +10,7 @@
 from torch import nn
 import transformer_engine_torch as tex
 from transformer_engine.pytorch.export import is_in_onnx_export_mode
+from transformer_engine import te_device_type
 
 THREADS_PER_WARP = 32
 THREADS_PER_BLOCK = 128
@@ -236,7 +237,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

@@ -22,6 +22,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,
     QKVFormat,
@@ -2038,7 +2039,7 @@ def _mask(cu_seqlens: torch.Tensor, max_seqlen: int) -> torch.Tensor:
         return (
             (positions.unsqueeze(0) >= seqlens.unsqueeze(1))
             .view(batch_size, 1, 1, max_seqlen)
-            .to(device="cuda")
+            .to(device=te_device_type())
         )
 
     attention_mask_q = _mask(cu_seqlens_q, max_seqlen_q)
@@ -2164,9 +2165,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 (
@@ -2262,7 +2265,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
 
@@ -2275,9 +2278,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:
@@ -2297,7 +2300,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"]
@@ -2312,7 +2317,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
@@ -2330,7 +2335,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)
@@ -2366,7 +2371,7 @@ def get_indices(max_seqlen: int, cu_seqlens: torch.Tensor) -> torch.Tensor:
         ordered,
         bs * max_seqlen,
     )
-    return indices.to(dtype=torch.int64, device="cuda").unsqueeze(1).unsqueeze(1)
+    return indices.to(dtype=torch.int64, device=te_device_type()).unsqueeze(1).unsqueeze(1)
 
 
 def get_full_cu_seqlens(

transformer_engine/pytorch/attention/inference.py

@@ -10,6 +10,7 @@
 
 import torch
 
+from transformer_engine import te_device_type
 from transformer_engine.pytorch.attention.custom_ops import (
     copy_to_kv_cache,
     QKV_FORMAT_VALUE,
@@ -631,7 +632,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):

transformer_engine/pytorch/attention/multi_head_attention.py

@@ -14,6 +14,7 @@
 from transformer_engine.pytorch.module.base import TransformerEngineBaseModule
 from transformer_engine.pytorch.module import LayerNormLinear, Linear, RMSNorm, LayerNorm
 from transformer_engine.pytorch.ops.basic.l2normalization import L2Normalization
+from transformer_engine import te_device_type
 from transformer_engine.pytorch.utils import (
     SplitAlongDim,
     divide,
@@ -141,7 +142,7 @@ class MultiheadAttention(torch.nn.Module):
                             whether to use interleaved rotary position embeddings.
     bias : bool, default = True
           if set to ``False``, the transformer layer will not learn any additive biases.
-    device : Union[torch.device, str], default = "cuda"
+    device : Union[torch.device, str], default=te_device_type()
           The device on which the parameters of the model will be allocated. It is the user's
           responsibility to ensure all parameters are moved to the GPU before running the
           forward pass.
@@ -285,7 +286,7 @@ def __init__(
         ub_bulk_wgrad: bool = False,
         bias: bool = True,
         normalization: str = "LayerNorm",
-        device: Union[torch.device, str] = "cuda",
+        device: Union[torch.device, str] = te_device_type(),
         qkv_format: str = "sbhd",
         name: str = None,
         qk_norm_type: Optional[str] = None,