diff --git a/transformer_engine/common/__init__.py b/transformer_engine/common/__init__.py
index e4a5f9256b8..ee50d14b7a6 100644
--- a/transformer_engine/common/__init__.py
+++ b/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
diff --git a/transformer_engine/debug/features/fake_quant.py b/transformer_engine/debug/features/fake_quant.py
index 48112abe4fe..570c0aadf73 100644
--- a/transformer_engine/debug/features/fake_quant.py
+++ b/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",
diff --git a/transformer_engine/debug/features/log_fp8_tensor_stats.py b/transformer_engine/debug/features/log_fp8_tensor_stats.py
index 96f1b644cfb..efc0d8793b3 100644
--- a/transformer_engine/debug/features/log_fp8_tensor_stats.py
+++ b/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)
diff --git a/transformer_engine/debug/features/per_tensor_scaling.py b/transformer_engine/debug/features/per_tensor_scaling.py
index e0fc44bd08f..f83b2795042 100644
--- a/transformer_engine/debug/features/per_tensor_scaling.py
+++ b/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,
diff --git a/transformer_engine/debug/features/utils/stats_buffer.py b/transformer_engine/debug/features/utils/stats_buffer.py
index ca7f22e2de5..51cc5a0c1f2 100644
--- a/transformer_engine/debug/features/utils/stats_buffer.py
+++ b/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
 
diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py
index aea10b50087..52f95340c44 100644
--- a/transformer_engine/pytorch/__init__.py
+++ b/transformer_engine/pytorch/__init__.py
@@ -9,6 +9,7 @@
 import functools
 
 import torch
+import transformer_engine
 
 from transformer_engine.common import load_framework_extension
 from transformer_engine.pytorch.torch_version import torch_version
@@ -16,6 +17,26 @@
 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
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py
index bccb4f1e26f..a04dd61474f 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py
+++ b/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 [
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py
index 3a153a3f994..376f8dc90be 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py
+++ b/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,
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
index e8e1396652d..91b917e4d01 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
+++ b/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,
             )
 
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py
index b9593b42d9b..2c8198881de 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py
+++ b/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,
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/softmax.py b/transformer_engine/pytorch/attention/dot_product_attention/softmax.py
index 6e8a14402f0..b55a960481d 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/softmax.py
+++ b/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():
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
index 710ed8f7fbf..67a768d70c4 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py
+++ b/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(
diff --git a/transformer_engine/pytorch/attention/inference.py b/transformer_engine/pytorch/attention/inference.py
index f98d1693a65..6746ceccaba 100644
--- a/transformer_engine/pytorch/attention/inference.py
+++ b/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):
diff --git a/transformer_engine/pytorch/attention/multi_head_attention.py b/transformer_engine/pytorch/attention/multi_head_attention.py
index e07e203ca7e..93ed14ad413 100644
--- a/transformer_engine/pytorch/attention/multi_head_attention.py
+++ b/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,
diff --git a/transformer_engine/pytorch/attention/rope.py b/transformer_engine/pytorch/attention/rope.py
index 77ad57ed8f7..18397f954b3 100644
--- a/transformer_engine/pytorch/attention/rope.py
+++ b/transformer_engine/pytorch/attention/rope.py
@@ -10,6 +10,7 @@
 
 import transformer_engine_torch as tex
 from transformer_engine.pytorch.cpp_extensions.fused_attn import QKVFormat
+from transformer_engine import te_device_type
 
 
 __all__ = ["RotaryPositionEmbedding", "apply_rotary_pos_emb", "apply_fused_qkv_rotary_pos_emb"]
@@ -76,7 +77,7 @@ def forward(self, max_seq_len: int, offset: int = 0):
         offset: int, default = 0
             Fixed offset for frequencies.
         """
-        with torch.autocast(enabled=False, device_type="cuda"):
+        with torch.autocast(enabled=False, device_type=te_device_type()):
             seq = (
                 torch.arange(max_seq_len, device=self.inv_freq.device, dtype=self.inv_freq.dtype)
                 + offset
diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py
index f2f8537cc23..bc670b7fec9 100644
--- a/transformer_engine/pytorch/cpp_extensions/gemm.py
+++ b/transformer_engine/pytorch/cpp_extensions/gemm.py
@@ -11,6 +11,7 @@
 import functools
 import torch
 import transformer_engine_torch as tex
+from transformer_engine import te_device_type
 from ..constants import MXFP8_BLOCK_SCALING_SIZE, NVFP4_BLOCK_SCALING_SIZE, TE_DType, DType
 from ..utils import ceil_div, get_cached_ones_tensor, get_sm_count, _empty_tensor
 
@@ -983,7 +984,8 @@ def general_grouped_gemm(
 
     if grad and use_bias:
         grad_bias = [
-            torch.empty(B[i].size(1), dtype=out[0].dtype, device="cuda") for i in range(num_gemms)
+            torch.empty(B[i].size(1), dtype=out[0].dtype, device=te_device_type())
+            for i in range(num_gemms)
         ]
     else:
         grad_bias = empty_tensors
diff --git a/transformer_engine/pytorch/cross_entropy.py b/transformer_engine/pytorch/cross_entropy.py
index 8ae3b853801..2318969963b 100644
--- a/transformer_engine/pytorch/cross_entropy.py
+++ b/transformer_engine/pytorch/cross_entropy.py
@@ -10,6 +10,7 @@
 import torch
 
 import transformer_engine.pytorch.triton.cross_entropy as triton_cross_entropy
+from transformer_engine import te_device_type
 from transformer_engine.pytorch.jit import no_torch_dynamo
 
 __all__ = ["parallel_cross_entropy"]
@@ -114,10 +115,10 @@ def _validate_inputs(
         raise ValueError(f"inp dimensions must be non-zero, but got shape {tuple(inp.shape)}")
     if inp.dtype not in (torch.bfloat16, torch.float32):
         raise TypeError(f"inp must have BF16 or FP32 dtype, but got {inp.dtype}")
-    if not inp.is_cuda:
-        raise ValueError("inp must be a CUDA tensor")
+    if inp.device.type != te_device_type():
+        raise ValueError(f"inp must be a {te_device_type()} tensor, but got {inp.device}")
     if target.device != inp.device:
-        raise ValueError("target must be on the same CUDA device as inp")
+        raise ValueError(f"target must be on the same {te_device_type()} device as inp")
     if target.dtype != torch.int64:
         raise TypeError(f"target must have torch.int64 dtype, but got {target.dtype}")
     if target.numel() != inp.shape[0] * inp.shape[1]:
diff --git a/transformer_engine/pytorch/custom_recipes/quantizer_factories.py b/transformer_engine/pytorch/custom_recipes/quantizer_factories.py
index 250ce8c9393..b6d22901dda 100644
--- a/transformer_engine/pytorch/custom_recipes/quantizer_factories.py
+++ b/transformer_engine/pytorch/custom_recipes/quantizer_factories.py
@@ -29,6 +29,7 @@
 
 import torch
 from transformer_engine.pytorch.quantization import QuantizerRole
+from transformer_engine import te_device_type
 from ..constants import DType
 
 
@@ -86,7 +87,7 @@ def current_scaling_factory(
 
     return Float8CurrentScalingQuantizer(
         fp8_dtype=fp8_dtype,
-        device=torch.device("cuda"),
+        device=torch.device(te_device_type()),
         force_pow_2_scales=False,
         amax_epsilon=0.0,
     )
diff --git a/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py b/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py
index 0d5fc0c4419..5cf573d61e4 100644
--- a/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py
+++ b/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py
@@ -72,6 +72,7 @@
 from typing import Optional
 
 from transformer_engine.pytorch.quantization import QuantizerRole
+from transformer_engine import te_device_type
 from ..constants import DType
 from .quantizer_factories import mxfp8_factory, nvfp4_factory
 
@@ -426,6 +427,6 @@ def nvfp4_linear_fp8_dpa_factory(
             is_dpa_boundary and "dpa_grad_input" in role.name
         )
         fp8_dtype = DType.kFloat8E5M2 if is_bwd_role else DType.kFloat8E4M3
-        return Float8CurrentScalingQuantizer(fp8_dtype=fp8_dtype, device="cuda")
+        return Float8CurrentScalingQuantizer(fp8_dtype=fp8_dtype, device=te_device_type())
 
     return nvfp4_factory(role)
diff --git a/transformer_engine/pytorch/custom_recipes/reference_current_scaling.py b/transformer_engine/pytorch/custom_recipes/reference_current_scaling.py
index dcb77bbbce1..6d016500c6a 100644
--- a/transformer_engine/pytorch/custom_recipes/reference_current_scaling.py
+++ b/transformer_engine/pytorch/custom_recipes/reference_current_scaling.py
@@ -13,6 +13,7 @@
 from transformer_engine.pytorch.custom_recipes import gemm
 from transformer_engine.pytorch.custom_recipes import reference_utils
 from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage, Quantizer
+from transformer_engine import te_device_type
 
 
 def current_scaling_ref_factory(role):
@@ -502,7 +503,7 @@ def make_empty(
 
         # Canonicalize tensor attributes
         if device is None:
-            device = torch.device("cuda")
+            device = torch.device(te_device_type())
 
         # Allocate quantized data
         qx = torch.empty(shape, dtype=self.dtype, device=device)
diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py
index e90b17e4d9f..9b59970ba88 100644
--- a/transformer_engine/pytorch/distributed.py
+++ b/transformer_engine/pytorch/distributed.py
@@ -31,6 +31,7 @@
 import transformer_engine_torch as tex
 
 from transformer_engine.pytorch.triton.pad import pad_columnwise_scale_inv
+from transformer_engine import te_device_type
 from .torch_version import torch_version
 from .utils import (
     is_non_tn_fp8_gemm_supported,
@@ -96,7 +97,7 @@ def is_graph_safe_rng_state(state: Union[torch.Tensor, torch.Generator]) -> bool
 
 
 def _get_cuda_rng_state(
-    device: Union[int, str, torch.device] = "cuda",
+    device: Union[int, str, torch.device] = te_device_type(),
     clone: bool = False,
     graph_safe: bool = True,
 ) -> torch.Tensor:
@@ -106,7 +107,7 @@ def _get_cuda_rng_state(
     if isinstance(device, str):
         device = torch.device(device)
     elif isinstance(device, int):
-        device = torch.device("cuda", device)
+        device = torch.device(te_device_type(), device)
     idx = device.index
     if idx is None:
         idx = torch.cuda.current_device()
@@ -128,11 +129,11 @@ def _set_cuda_rng_state(
     """Sets the random number generator state of the current GPU."""
 
     if device == -1:
-        device = torch.device("cuda")
+        device = torch.device(te_device_type())
     elif isinstance(device, str):
         device = torch.device(device)
     elif isinstance(device, int):
-        device = torch.device("cuda", device)
+        device = torch.device(te_device_type(), device)
 
     def cb() -> None:
         idx = device.index
@@ -294,10 +295,10 @@ def _get_active_autocast_contexts():
     autocast_cached = torch.is_autocast_cache_enabled()
 
     if torch_version() >= (2, 4, 0):
-        gpu_autocast_enabled = torch.is_autocast_enabled("cuda")
-        gpu_autocast_dtype = torch.get_autocast_dtype("cuda")
+        gpu_autocast_enabled = torch.is_autocast_enabled(te_device_type())
+        gpu_autocast_dtype = torch.get_autocast_dtype(te_device_type())
         gpu_autocast_ctx = torch.amp.autocast(
-            "cuda",
+            te_device_type(),
             enabled=gpu_autocast_enabled,
             dtype=gpu_autocast_dtype,
             cache_enabled=autocast_cached,
@@ -1047,7 +1048,7 @@ def _all_gather_fp8(
     out: Float8TensorStorage
     if quantizer is not None:
         dtype = torch.float32
-        device = "cuda"
+        device = te_device_type()
         if isinstance(inp, Float8Tensor):
             dtype = inp.dtype
             device = inp.device
@@ -1987,7 +1988,7 @@ def symm_mem_alloc(
     be a shared static buffer. ``backend`` selects the symm-mem backend (default NCCL; for the pool it
     is captured at pool creation)."""
     if device is None:
-        device = torch.device("cuda", torch.cuda.current_device())
+        device = torch.device(te_device_type(), torch.cuda.current_device())
     if not HAS_TORCH_SYMMETRIC:
         raise RuntimeError(
             "torch.distributed._symmetric_memory is unavailable; symm_mem_alloc "
diff --git a/transformer_engine/pytorch/dynamo/tensor_spec.py b/transformer_engine/pytorch/dynamo/tensor_spec.py
index d96c49ca7a5..e5afa0fdd0b 100644
--- a/transformer_engine/pytorch/dynamo/tensor_spec.py
+++ b/transformer_engine/pytorch/dynamo/tensor_spec.py
@@ -11,6 +11,7 @@
 
 import torch
 from torch._prims_common import make_contiguous_strides_for
+from transformer_engine import te_device_type
 
 
 @dataclass
@@ -108,7 +109,7 @@ def create_inner_tensors(self) -> List[torch.Tensor]:
         Under ``register_fake`` the ``torch.empty`` calls produce ``FakeTensor``s;
         ``requires_grad`` is left default (managed by ``register_autograd``).
         """
-        device = self.device if self.device is not None else torch.device("cuda")
+        device = self.device if self.device is not None else torch.device(te_device_type())
         if self.quantizer is None:
             return [torch.empty(tuple(self.shape), dtype=self.dtype, device=device)]
         inner = self.quantizer.alloc_tensors(tuple(self.shape), device=device)
@@ -139,7 +140,7 @@ def create_tensor(self) -> torch.Tensor:
         inner tensors via :meth:`assemble`.
         """
         if self.quantizer is None:
-            device = self.device if self.device is not None else torch.device("cuda")
+            device = self.device if self.device is not None else torch.device(te_device_type())
             return torch.empty(
                 tuple(self.shape),
                 dtype=self.dtype,
diff --git a/transformer_engine/pytorch/ep.py b/transformer_engine/pytorch/ep.py
index 9ed10fb4721..9709f4c7ed7 100644
--- a/transformer_engine/pytorch/ep.py
+++ b/transformer_engine/pytorch/ep.py
@@ -15,6 +15,7 @@
 
 import transformer_engine_torch as tex
 
+from transformer_engine import te_device_type
 from .constants import DType
 from .cpu_offload import mark_not_offload
 from .distributed import symm_mem_alloc, release_symm_mem_pool
@@ -349,7 +350,7 @@ def __init__(
         if not _BOOTSTRAPPED:
             raise RuntimeError("EpBuffer requires ep_bootstrap() to be called first.")
         if device is None:
-            device = torch.device("cuda", torch.cuda.current_device())
+            device = torch.device(te_device_type(), torch.cuda.current_device())
         alignment = int(alignment)
         if alignment > 1 and (alignment & (alignment - 1)) != 0:
             raise ValueError(f"alignment must be 0, 1, or a power of two (got {alignment}).")
diff --git a/transformer_engine/pytorch/jit.py b/transformer_engine/pytorch/jit.py
index 8bd2f68a070..8052f65bc2b 100644
--- a/transformer_engine/pytorch/jit.py
+++ b/transformer_engine/pytorch/jit.py
@@ -10,6 +10,7 @@
 from typing import Callable, Optional, Tuple
 import torch
 
+from .. import te_device_type
 from .torch_version import torch_version
 from .export import is_in_onnx_export_mode
 from .utils import gpu_autocast_ctx
@@ -337,9 +338,17 @@ def warmup_jit_bias_dropout_add(
     # Save cuda RNG state to ensure warmup does not affect reproducibility.
     rng_state = torch.cuda.get_rng_state()
 
-    inp = torch.rand((seq_length, micro_batch_size, hidden_size), dtype=dtype, device="cuda")
-    residual = torch.rand((seq_length, micro_batch_size, hidden_size), dtype=dtype, device="cuda")
-    bias = torch.rand((hidden_size), dtype=dtype, device="cuda")
+    inp = torch.rand(
+        (seq_length, micro_batch_size, hidden_size),
+        dtype=dtype,
+        device=te_device_type(),
+    )
+    residual = torch.rand(
+        (seq_length, micro_batch_size, hidden_size),
+        dtype=dtype,
+        device=te_device_type(),
+    )
+    bias = torch.rand((hidden_size), dtype=dtype, device=te_device_type())
     dropout_rate = 0.1
     # Warmup JIT fusions with the input grad_enable state of both forward
     # prop and recomputation
@@ -374,11 +383,11 @@ def warmup_jit_bias_gelu(
     # Save cuda RNG state to ensure warmup does not affect reproducibility.
     rng_state = torch.cuda.get_rng_state()
 
-    bias = torch.rand(ffn_hidden_size_per_partition, dtype=dtype, device="cuda")
+    bias = torch.rand(ffn_hidden_size_per_partition, dtype=dtype, device=te_device_type())
     inp = torch.rand(
         (seq_length * micro_batch_size, ffn_hidden_size_per_partition),
         dtype=dtype,
-        device="cuda",
+        device=te_device_type(),
     )
     # Warmup JIT fusions with the input grad_enable state of both forward
     # prop and recomputation
@@ -412,7 +421,7 @@ def warmup_jit_l2normalization(
     inp = torch.rand(
         (seq_length * micro_batch_size, hidden_size),
         dtype=dtype,
-        device="cuda",
+        device=te_device_type(),
     )
     eps = 1e-6
     # Warmup JIT fusions with the input grad_enable state of both forward
diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py
index 5aea8feea1c..2e998761154 100644
--- a/transformer_engine/pytorch/module/base.py
+++ b/transformer_engine/pytorch/module/base.py
@@ -21,6 +21,7 @@
 
 import transformer_engine_torch as tex
 
+from transformer_engine import te_device_type
 from ._common import _ParameterInitMeta, noop_cat, sum_bias_grad
 from .._extra_state import (
     extra_state_pickle_advisory,
@@ -63,6 +64,7 @@
 from ..utils import (
     is_non_tn_fp8_gemm_supported,
     torch_get_autocast_gpu_dtype,
+    torch_is_autocast_gpu_enabled,
     get_device_compute_capability,
     get_nvtx_range_context,
     nvtx_range_push,
@@ -146,7 +148,7 @@ def get_dummy_wgrad(
 
     # The dummy goes on the current device (the one the backward pass runs on), and the cache
     # is keyed by it so that a dummy is never returned for a weight on another device.
-    device = torch.device("cuda", torch.cuda.current_device())
+    device = torch.device(te_device_type(), torch.cuda.current_device())
     if zero_tensor and not zero:
         return torch._efficientzerotensor(shape, dtype=dtype, device=device)
     key = (*shape, dtype, device)
@@ -1554,7 +1556,7 @@ def set_extra_state(self, state: torch.Tensor) -> None:
             if not unsafe_pickle_extra_state_enabled():
                 raise RuntimeError(extra_state_pickle_advisory(context))
             state.seek(0)
-            state = torch.load(state, map_location="cuda", weights_only=False)
+            state = torch.load(state, map_location=te_device_type(), weights_only=False)
         else:
             raise RuntimeError("Unsupported checkpoint format.")
 
@@ -1598,7 +1600,7 @@ def copy_tensor(src: torch.Tensor, dst: torch.Tensor) -> None:
     def set_activation_dtype(self, inp: torch.Tensor) -> None:
         """Get activation data type for AMP."""
         # Native AMP (`torch.autocast`) gets highest priority
-        if torch.is_autocast_enabled():
+        if torch_is_autocast_gpu_enabled():
             self.fast_setattr("activation_dtype", torch_get_autocast_gpu_dtype())
             return
         dtype = inp.dtype
@@ -1718,9 +1720,9 @@ def prepare_forward(
             delayed_scaling_recipe = _has_delayed_scaling_state(self.fp8_meta)
             FP8GlobalStateManager.get_old_fp8_meta_tensors_for_recompute(self.fp8_meta)
         else:
-            if not inp.is_cuda:
+            if inp.device.type != te_device_type():
                 raise RuntimeError(
-                    f"TransformerEngine needs CUDA. Got input on device: {inp.device}"
+                    f"TransformerEngine needs {te_device_type()}. Got input on device: {inp.device}"
                 )
 
             if self.tp_size > 1:
@@ -1934,7 +1936,7 @@ def reset_parameters(self, defer_init: Optional[bool] = False) -> None:
             param = param._local_tensor if is_dtensor else param
             # Ensure parameter is on a real device
             if param.device == torch.device("meta"):
-                param = torch.empty_like(param, device="cuda")
+                param = torch.empty_like(param, device=te_device_type())
             # Initialize the parameter values on device
             init_fn = self.param_init_meta[name].init_fn
             get_rng_state_tracker = self.param_init_meta[name].get_rng_state_tracker
diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py
index 283deecfdd2..d3fea572398 100644
--- a/transformer_engine/pytorch/module/grouped_linear.py
+++ b/transformer_engine/pytorch/module/grouped_linear.py
@@ -17,6 +17,7 @@
 import transformer_engine_torch as tex
 
 from transformer_engine.common.recipe import Recipe, Format as RecipeFormat
+from transformer_engine import te_device_type
 from transformer_engine.pytorch.tensor.grouped_tensor import (
     GroupedTensor,
     GroupedTensorStorage,
@@ -1629,7 +1630,7 @@ class GroupedLinear(TransformerEngineBaseModule):
                  used to get the random number generator state tracker for initializing weights.
     rng_tracker_name : str, default = None
                  the param passed to get_rng_state_tracker to get the specific rng tracker.
-    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.
@@ -1707,7 +1708,7 @@ def __init__(
         return_bias: bool = False,
         params_dtype: Optional[torch.dtype] = None,
         parallel_mode: Optional[str] = None,
-        device: Union[torch.device, str] = "cuda",
+        device: Union[torch.device, str] = te_device_type(),
         ub_overlap_rs: bool = False,
         ub_overlap_ag: bool = False,
         ub_name: Optional[str] = None,
@@ -2416,7 +2417,7 @@ def forward(
                     " single_grouped_bias to allow the split-quantize fallback."
                 )
             if use_grouped_tensor_path:
-                if m_splits.device.type != "cuda":
+                if m_splits.device.type != te_device_type():
                     raise ValueError(
                         "The native grouped_tensor path requires CUDA m_splits. Pass a CUDA int64 "
                         "tensor, or set use_grouped_tensor=False."
diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py
index 444b66f2eb1..03e6cf50de8 100644
--- a/transformer_engine/pytorch/module/layernorm_linear.py
+++ b/transformer_engine/pytorch/module/layernorm_linear.py
@@ -17,6 +17,7 @@
 from transformer_engine.common.recipe import Recipe
 from transformer_engine.pytorch.torch_version import torch_version
 from transformer_engine.pytorch.tensor.utils import clear_columnwise_cache, is_custom
+from transformer_engine import te_device_type
 from .base import (
     fill_userbuffers_buffer_for_all_gather,
     get_cublasmp_all_gather_output,
@@ -1600,7 +1601,7 @@ class LayerNormLinear(TransformerEngineBaseModule):
                          .. math::
                             y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \varepsilon}} *
                             (1 + \gamma) + \beta
-    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.
@@ -1674,7 +1675,7 @@ def __init__(
         return_layernorm_output_gathered: bool = False,
         parameters_split: Optional[Union[Tuple[str, ...], Dict[str, int]]] = None,
         zero_centered_gamma: bool = False,
-        device: Union[torch.device, str] = "cuda",
+        device: Union[torch.device, str] = te_device_type(),
         ub_overlap_ag: bool = False,
         ub_overlap_rs: bool = False,
         ub_overlap_rs_dgrad: bool = False,
diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py
index 10c404bb5d5..92ae0c9ac23 100644
--- a/transformer_engine/pytorch/module/layernorm_mlp.py
+++ b/transformer_engine/pytorch/module/layernorm_mlp.py
@@ -18,6 +18,7 @@
 from transformer_engine.common.recipe import Recipe
 from transformer_engine.pytorch.torch_version import torch_version
 from transformer_engine.pytorch.tensor.utils import clear_columnwise_cache, is_custom
+from transformer_engine import te_device_type
 from .base import (
     fill_userbuffers_buffer_for_all_gather,
     get_cublasmp_all_gather_output,
@@ -1798,7 +1799,7 @@ def fc2_wgrad_gemm(
         reduce_scatter_out = None
         if args.ub_overlap_rs_dgrad:
             reduce_scatter_out = torch.empty(
-                fc1_dgrad_shape, dtype=args.activation_dtype, device="cuda"
+                fc1_dgrad_shape, dtype=args.activation_dtype, device=te_device_type()
             )
         if ub_bulk_wgrad:
             gemm_out = ub_obj_fc1_wgrad.get_buffer(
@@ -1896,7 +1897,7 @@ def fc2_wgrad_gemm(
             reduce_scatter_out = None
             if ub_bulk_wgrad and ub_obj_fc1_wgrad.is_fp8_ubuf():
                 reduce_scatter_out = torch.empty(
-                    fc1_dgrad_shape, dtype=args.activation_dtype, device="cuda"
+                    fc1_dgrad_shape, dtype=args.activation_dtype, device=te_device_type()
                 )
 
             # Arguments to include in wgrad GEMM closure
@@ -2222,7 +2223,7 @@ class LayerNormMLP(TransformerEngineBaseModule):
                          .. math::
                             y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \varepsilon}} *
                             (1 + \gamma) + \beta
-    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.
@@ -2311,7 +2312,7 @@ def __init__(
         micro_batch_size: Optional[int] = None,
         set_parallel_mode: bool = False,
         zero_centered_gamma: bool = False,
-        device: Union[torch.device, str] = "cuda",
+        device: Union[torch.device, str] = te_device_type(),
         ub_overlap_ag: bool = False,
         name: Optional[str] = None,
         ub_overlap_rs: bool = False,
diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
index 87d76ec02cf..ba43c9fcc1c 100644
--- a/transformer_engine/pytorch/module/linear.py
+++ b/transformer_engine/pytorch/module/linear.py
@@ -16,6 +16,7 @@
 from transformer_engine.common.recipe import Recipe
 from transformer_engine.pytorch.torch_version import torch_version
 
+from transformer_engine import te_device_type
 from .base import (
     fill_userbuffers_buffer_for_all_gather,
     get_cublasmp_all_gather_output,
@@ -1958,7 +1959,7 @@ class Linear(TransformerEngineBaseModule):
                       values as split sizes along dim 0. The resulting parameters will have
                       names that end in ``_weight`` or ``_bias``, so trailing underscores are
                       stripped from any provided names.
-    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.
@@ -2034,7 +2035,7 @@ def __init__(
         params_dtype: Optional[torch.dtype] = None,
         parallel_mode: Optional[str] = None,
         parameters_split: Optional[Union[Tuple[str, ...], Dict[str, int]]] = None,
-        device: Union[torch.device, str] = "cuda",
+        device: Union[torch.device, str] = te_device_type(),
         ub_overlap_ag: bool = False,
         ub_overlap_rs: bool = False,
         ub_overlap_rs_dgrad: bool = False,
@@ -2345,7 +2346,7 @@ def _prealloc_cublas_workspace(self) -> None:
         from torch._subclasses.fake_tensor import is_fake
 
         weight = getattr(self, self.weight_names[0], None)
-        if weight is None or weight.device.type != "cuda":
+        if weight is None or weight.device.type != te_device_type():
             return
         if is_fake(weight) or detect_fake_mode():
             return
diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py
index cf4e66023a6..e86b0b5a778 100644
--- a/transformer_engine/pytorch/ops/_common.py
+++ b/transformer_engine/pytorch/ops/_common.py
@@ -11,6 +11,7 @@
 import torch
 
 from transformer_engine_torch import FP8TensorMeta
+from transformer_engine import te_device_type
 from ..dynamo.tensor_spec import TensorSpec
 from ..ep import (
     EpBuffer,
@@ -179,7 +180,7 @@ def maybe_dequantize(
 
 def maybe_autocast_dtype(
     *,
-    device_type: str = "cuda",
+    device_type: str = te_device_type(),
     default_dtype: Optional[torch.dtype] = None,
 ) -> torch.dtype:
     """Get autocast dtype if enabled"""
diff --git a/transformer_engine/pytorch/ops/basic/activation.py b/transformer_engine/pytorch/ops/basic/activation.py
index ffcb72a9554..21a376ca3d1 100644
--- a/transformer_engine/pytorch/ops/basic/activation.py
+++ b/transformer_engine/pytorch/ops/basic/activation.py
@@ -18,7 +18,7 @@
 from ...tensor.float8_tensor import Float8CurrentScalingQuantizer, Quantizer
 from ...utils import _compile_safe_warn, clear_tensor_data
 from ..op import BasicOperation, OperationContext
-from .._common import maybe_dequantize
+from .._common import maybe_autocast_dtype, maybe_dequantize
 
 __all__ = [
     "GELU",
@@ -94,11 +94,7 @@ def op_forward(
     ) -> torch.Tensor:
 
         # Compute dtype
-        dtype: torch.dtype
-        if torch.is_autocast_enabled():
-            dtype = torch.get_autocast_dtype("cuda")
-        else:
-            dtype = input_.dtype
+        dtype: torch.dtype = maybe_autocast_dtype(default_dtype=input_.dtype)
         if dtype not in (torch.float32, torch.float16, torch.bfloat16):
             raise RuntimeError(f"Unsupported dtype ({dtype})")
 
@@ -412,12 +408,8 @@ def fuser_forward(
                 "in the fused grouped MLP path."
             )
 
-        if torch.is_autocast_enabled():
-            dtype = torch.get_autocast_dtype("cuda")
-            extra_input_dtype = dtype
-        else:
-            dtype = input_.dtype
-            extra_input_dtype = extra_input.dtype
+        dtype = maybe_autocast_dtype(default_dtype=input_.dtype)
+        extra_input_dtype = maybe_autocast_dtype(default_dtype=extra_input.dtype)
 
         x = maybe_dequantize(input_.contiguous(), dtype)
         scales = maybe_dequantize(extra_input.contiguous(), extra_input_dtype)
diff --git a/transformer_engine/pytorch/ops/basic/basic_linear.py b/transformer_engine/pytorch/ops/basic/basic_linear.py
index df8e13604a9..e60662226d3 100644
--- a/transformer_engine/pytorch/ops/basic/basic_linear.py
+++ b/transformer_engine/pytorch/ops/basic/basic_linear.py
@@ -43,6 +43,7 @@
     get_dummy_wgrads_for_params,
     get_main_grad_from_param,
     is_quantized_tensor,
+    maybe_autocast_dtype,
     maybe_dequantize,
 )
 
@@ -1108,11 +1109,7 @@ def pack_forward_args(
         return BasicLinearFwdArgs(
             input_=input_,
             weight=self.weight,
-            dtype=(
-                torch.get_autocast_dtype("cuda")
-                if torch.is_autocast_enabled()
-                else self.weight.dtype
-            ),
+            dtype=maybe_autocast_dtype(default_dtype=self.weight.dtype),
             input_requires_grad=basic_op_ctxs[0].requires_grad,
             weight_requires_grad=basic_op_ctxs[0].requires_grad and self.weight.requires_grad,
             with_quantized_compute=with_quantized_compute,
diff --git a/transformer_engine/pytorch/ops/basic/bias.py b/transformer_engine/pytorch/ops/basic/bias.py
index 0dab1e9ed0b..0a26e0460b9 100644
--- a/transformer_engine/pytorch/ops/basic/bias.py
+++ b/transformer_engine/pytorch/ops/basic/bias.py
@@ -11,6 +11,7 @@
 import torch
 
 import transformer_engine_torch as tex
+from transformer_engine import te_device_type
 from ...dynamo import TensorOrQuantized, TensorSpec
 from ...dynamo.fake import fake_bgrad
 from ...quantization import FP8GlobalStateManager
@@ -115,7 +116,7 @@ def reset_parameters(self) -> None:
 
         # Make sure parameter is initialized
         bias = self.bias
-        if bias.device.type != "cuda":
+        if bias.device.type != te_device_type():
             bias = torch.empty_like(bias, device=self.device)
         else:
             bias = bias.to(device=self.device)
diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py
index 91f627148bf..4b68426a0ab 100644
--- a/transformer_engine/pytorch/ops/basic/grouped_linear.py
+++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py
@@ -49,6 +49,7 @@
     get_dummy_wgrads_for_params,
     get_main_grad_from_param,
     is_quantized_tensor,
+    maybe_autocast_dtype,
     maybe_dequantize,
     validate_or_alloc_output,
     view_main_grad_as_grouped_buffer,
@@ -1054,10 +1055,7 @@ def fuser_forward(
                 weight_quantizers[group_idx] = self.get_quantizer("forward", 2 * group_idx + 1)
 
         # Get autocast dtype if needed
-        if torch.is_autocast_enabled():
-            dtype = torch.get_autocast_dtype("cuda")
-        else:
-            dtype = weight_param.dtype
+        dtype = maybe_autocast_dtype(default_dtype=weight_param.dtype)
 
         # Extract split sizes from extra input. Keep on GPU for graph safety.
         split_sizes = basic_op_extra_inputs[0][0]
@@ -1208,10 +1206,7 @@ def fuser_forward_save_ctx(
         # ``split_sizes``, offset metadata, and related tensors are routed
         # through ``save_for_backward`` (see ``_fuser_forward_split_quantize``
         # and ``_fuser_forward_grouped_tensor`` for the saved-tensor layout).
-        if torch.is_autocast_enabled():
-            ctx.dtype = torch.get_autocast_dtype("cuda")
-        else:
-            ctx.dtype = weight_param.dtype
+        ctx.dtype = maybe_autocast_dtype(default_dtype=weight_param.dtype)
         ctx.input_requires_grad = requires_grad[0]
         ctx.weight_requires_grad = requires_grad[0] and weight_param.requires_grad
         # Caller-provided backward grad-input buffer.
diff --git a/transformer_engine/pytorch/ops/basic/swiglu.py b/transformer_engine/pytorch/ops/basic/swiglu.py
index 8009de9b4ae..71ff010ae53 100644
--- a/transformer_engine/pytorch/ops/basic/swiglu.py
+++ b/transformer_engine/pytorch/ops/basic/swiglu.py
@@ -17,7 +17,7 @@
 from ...tensor import Float8CurrentScalingQuantizer, Quantizer
 from ...utils import clear_tensor_data
 from ..op import BasicOperation, OperationContext
-from .._common import maybe_dequantize
+from .._common import maybe_autocast_dtype, maybe_dequantize
 
 __all__ = [
     "SwiGLU",
@@ -114,11 +114,7 @@ def op_forward(
     ) -> torch.Tensor:
 
         # Compute dtype
-        dtype: torch.dtype
-        if torch.is_autocast_enabled():
-            dtype = torch.get_autocast_dtype("cuda")
-        else:
-            dtype = input_.dtype
+        dtype: torch.dtype = maybe_autocast_dtype(default_dtype=input_.dtype)
         if dtype not in (torch.float32, torch.float16, torch.bfloat16):
             raise RuntimeError(f"Unsupported dtype ({dtype})")
 
@@ -361,11 +357,7 @@ def op_forward(
     ) -> torch.Tensor:
 
         # Compute dtype
-        dtype: torch.dtype
-        if torch.is_autocast_enabled():
-            dtype = torch.get_autocast_dtype("cuda")
-        else:
-            dtype = input_.dtype
+        dtype: torch.dtype = maybe_autocast_dtype(default_dtype=input_.dtype)
         if dtype not in (torch.float32, torch.float16, torch.bfloat16):
             raise RuntimeError(f"Unsupported dtype ({dtype})")
 
@@ -528,12 +520,8 @@ def fuser_forward(
 
         extra_input = basic_op_extra_inputs[0][0]
         # Determine compute dtype
-        if torch.is_autocast_enabled():
-            dtype = torch.get_autocast_dtype("cuda")
-            extra_input_dtype = dtype
-        else:
-            dtype = input_.dtype
-            extra_input_dtype = extra_input.dtype
+        dtype = maybe_autocast_dtype(default_dtype=input_.dtype)
+        extra_input_dtype = maybe_autocast_dtype(default_dtype=extra_input.dtype)
 
         # Prepare the activation input in the compute dtype.
         input_ = maybe_dequantize(input_, dtype)
diff --git a/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py b/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py
index 28586360f59..7f7a6bd97e3 100644
--- a/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py
+++ b/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py
@@ -13,6 +13,7 @@
 from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload
 from ...quantization import FP8GlobalStateManager
 from ...tensor import Quantizer
+from .._common import maybe_autocast_dtype
 from ..basic import AddExtraInput, BasicLinear, Bias
 from ..op import FusedOperation, FusibleOperation, OperationContext
 
@@ -92,10 +93,7 @@ def fuser_forward(
             backward_override = None
 
         # Get autocast dtype if needed
-        if torch.is_autocast_enabled():
-            dtype = torch.get_autocast_dtype("cuda")
-        else:
-            dtype = linear_op.weight.dtype
+        dtype = maybe_autocast_dtype(default_dtype=linear_op.weight.dtype)
 
         # Linear forward
         output = basic_op_extra_inputs[self._op_idxs["add"]][0]
diff --git a/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py b/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py
index 277263e0ecb..da78bec23b7 100644
--- a/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py
+++ b/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py
@@ -13,6 +13,7 @@
 from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload
 from ...quantization import FP8GlobalStateManager
 from ...tensor import Quantizer
+from .._common import maybe_autocast_dtype
 from ..basic import AddExtraInput, BasicLinear, ConstantScale
 from ..op import (
     FusedOperation,
@@ -74,10 +75,7 @@ def fuser_forward(
         extra_input = basic_op_extra_inputs[2][0]
 
         # Get autocast dtype if needed
-        if torch.is_autocast_enabled():
-            dtype = torch.get_autocast_dtype("cuda")
-        else:
-            dtype = linear_op.weight.dtype
+        dtype = maybe_autocast_dtype(default_dtype=linear_op.weight.dtype)
 
         # Linear forward
         output, x_local, w = BasicLinear._functional_forward(
diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py
index eeb5f43335f..0f555536f96 100644
--- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py
+++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py
@@ -59,6 +59,7 @@
     get_dummy_wgrads_for_params,
     get_main_grad_from_param,
     is_quantized_tensor,
+    maybe_autocast_dtype,
     maybe_dequantize,
     validate_or_alloc_output,
     view_main_grad_as_grouped_buffer,
@@ -1343,10 +1344,7 @@ def fuser_forward(
             assert fc2_op.weight0.is_routed_expert and fc2_op.weight0.weight_list is not None
 
         device = fc1_weight_param.device
-        if torch.is_autocast_enabled():
-            dtype = torch.get_autocast_dtype("cuda")
-        else:
-            dtype = fc1_weight_param.dtype
+        dtype = maybe_autocast_dtype(default_dtype=fc1_weight_param.dtype)
 
         # Check which grads are required
         requires_grad = any(ctx.requires_grad for ctx in basic_op_ctxs)
diff --git a/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py b/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py
index 37c3dee9a16..39b2b6fe136 100644
--- a/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py
+++ b/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py
@@ -11,6 +11,7 @@
 import torch
 
 from transformer_engine_torch import CommOverlapType, bulk_overlap_ag_with_external_gemm
+from transformer_engine import te_device_type
 from ...cpp_extensions import general_gemm
 from ...distributed import gather_along_first_dim, get_distributed_world_size
 from ...module.base import (
@@ -182,7 +183,7 @@ def _functional_backward(
             else:
                 device = grad_output.device
         device = canonicalize_device(device)
-        if device.type != "cuda":
+        if device.type != te_device_type():
             raise ValueError(f"Only CUDA devices are supported (got {device})")
 
         # Check datatype
diff --git a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py
index cb5f8a16e61..9bc2cfc1e75 100644
--- a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py
+++ b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py
@@ -11,6 +11,7 @@
 import torch
 
 from transformer_engine_torch import CommOverlapType
+from transformer_engine import te_device_type
 from ...cpp_extensions import general_gemm
 from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload
 from ...distributed import get_distributed_world_size
@@ -23,7 +24,7 @@
 from ...quantized_tensor import Quantizer
 from ...tensor.float8_tensor import Float8Quantizer, Float8CurrentScalingQuantizer
 from ...tensor.storage.float8_tensor_storage import Float8TensorStorage
-from .._common import maybe_dequantize, is_quantized_tensor
+from .._common import maybe_autocast_dtype, maybe_dequantize, is_quantized_tensor
 from ..basic import BasicLinear, Bias, ReduceScatter
 from ..op import (
     FusedOperation,
@@ -155,7 +156,7 @@ def _functional_forward(
         """
 
         # Check device
-        if device.type != "cuda":
+        if device.type != te_device_type():
             raise ValueError(f"Only CUDA devices are supported (got {device})")
 
         # Check datatype
@@ -321,10 +322,7 @@ def fuser_forward(
                 )
 
         # Get autocast dtype if needed
-        if torch.is_autocast_enabled():
-            dtype = torch.get_autocast_dtype("cuda")
-        else:
-            dtype = linear_op.weight.dtype
+        dtype = maybe_autocast_dtype(default_dtype=linear_op.weight.dtype)
 
         # Userbuffers options
         if linear_op._userbuffers_options is None:
diff --git a/transformer_engine/pytorch/optimizers/fused_adam.py b/transformer_engine/pytorch/optimizers/fused_adam.py
index b4b7f707a83..2cde7780af0 100644
--- a/transformer_engine/pytorch/optimizers/fused_adam.py
+++ b/transformer_engine/pytorch/optimizers/fused_adam.py
@@ -15,6 +15,7 @@
 import transformer_engine_torch as tex
 from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor, Float8Quantizer
 from transformer_engine.pytorch.quantized_tensor import QuantizedTensor
+from transformer_engine import te_device_type
 from ..constants import DType
 from .multi_tensor_apply import multi_tensor_applier
 
@@ -190,7 +191,7 @@ def __init__(
             self._step_supports_amp_scaling = True
 
         # Skip buffer
-        self._dummy_overflow_buf = torch.tensor([0], dtype=torch.int, device="cuda")
+        self._dummy_overflow_buf = torch.tensor([0], dtype=torch.int, device=te_device_type())
         self.multi_tensor_adam = tex.multi_tensor_adam
         self.multi_tensor_adam_param_remainder = tex.multi_tensor_adam_param_remainder
         self.multi_tensor_adam_fp8 = tex.multi_tensor_adam_fp8
diff --git a/transformer_engine/pytorch/permutation.py b/transformer_engine/pytorch/permutation.py
index 5e6bdb4dfcd..50880bcd7f7 100644
--- a/transformer_engine/pytorch/permutation.py
+++ b/transformer_engine/pytorch/permutation.py
@@ -7,6 +7,7 @@
 from typing import Optional, Tuple
 import torch
 import transformer_engine_torch as tex
+from transformer_engine import te_device_type
 import transformer_engine.pytorch.triton.permutation as triton_permutation
 from transformer_engine.pytorch.constants import TE_DType
 from transformer_engine.pytorch.quantized_tensor import (
@@ -44,10 +45,14 @@ def moe_permute_index_map_forward(
     if not inp.numel():
         return inp.clone(), torch.empty(0, dtype=torch.int32, device=inp.device)
 
-    if not inp.is_cuda:
-        raise ValueError(f"inp must be a CUDA tensor, but got tensor on {inp.device}.")
-    if not index.is_cuda:
-        raise ValueError(f"index must be a CUDA tensor, but got tensor on {index.device}.")
+    if inp.device.type != te_device_type():
+        raise ValueError(
+            f"inp must be a {te_device_type()} tensor, but got tensor on {inp.device}."
+        )
+    if index.device.type != te_device_type():
+        raise ValueError(
+            f"index must be a {te_device_type()} tensor, but got tensor on {index.device}."
+        )
     if inp.size(0) != index.size(0):
         raise ValueError(
             f"Permute not possible: inp.size(0) ({inp.size(0)}) must match "
@@ -299,19 +304,25 @@ def moe_permute_mask_map_forward(
         )
         return inp.clone(), row_id_map, permuted_probs
 
-    if not inp.is_cuda:
-        raise ValueError(f"inp must be a CUDA tensor, but got tensor on {inp.device}.")
-    if not routing_map.is_cuda:
+    if inp.device.type != te_device_type():
         raise ValueError(
-            f"routing_map must be a CUDA tensor, but got tensor on {routing_map.device}."
+            f"inp must be a {te_device_type()} tensor, but got tensor on {inp.device}."
+        )
+    if routing_map.device.type != te_device_type():
+        raise ValueError(
+            f"routing_map must be a {te_device_type()} tensor, but got tensor on"
+            f" {routing_map.device}."
         )
     if probs is not None:
-        if not probs.is_cuda:
-            raise ValueError(f"probs must be a CUDA tensor, but got tensor on {probs.device}.")
+        if probs.device.type != te_device_type():
+            raise ValueError(
+                f"probs must be a {te_device_type()} tensor, but got tensor on {probs.device}."
+            )
     if pad_offsets is not None:
-        if not pad_offsets.is_cuda:
+        if pad_offsets.device.type != te_device_type():
             raise ValueError(
-                f"pad_offsets must be a CUDA tensor, but got tensor on {pad_offsets.device}."
+                f"pad_offsets must be a {te_device_type()} tensor, but got tensor on"
+                f" {pad_offsets.device}."
             )
     if inp.size(0) != routing_map.size(0):
         raise ValueError(
diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py
index c167e5bb243..06d02ece44e 100644
--- a/transformer_engine/pytorch/quantization.py
+++ b/transformer_engine/pytorch/quantization.py
@@ -16,6 +16,8 @@
 
 import torch
 import transformer_engine_torch as tex
+
+from transformer_engine import te_device_type
 from transformer_engine.common.recipe import (
     Recipe,
     DelayedScaling,
@@ -480,7 +482,7 @@ def set_skip_fp8_weight_update_tensor(cls, skip: bool) -> None:
         """Set the skip fp8 weight update tensor"""
         if cls.quantization_state.skip_fp8_weight_update_tensor is None:
             cls.quantization_state.skip_fp8_weight_update_tensor = torch.empty(
-                1, dtype=torch.float32, device="cuda"
+                1, dtype=torch.float32, device=te_device_type()
             )
         cls.quantization_state.skip_fp8_weight_update_tensor.fill_(skip)
 
@@ -1488,7 +1490,7 @@ def __init__(
 
         # Allocate buffers
         if device is None:
-            device = torch.device("cuda")
+            device = torch.device(te_device_type())
         self.scale = torch.ones(num_quantizers, dtype=torch.float32, device=device)
         self.amax_history = torch.zeros(
             recipe.amax_history_len,
@@ -1537,7 +1539,7 @@ def __init__(
 
         # Allocate buffers
         if device is None:
-            device = torch.device("cuda")
+            device = torch.device(te_device_type())
         self.device = device
 
     def make_quantizers(self) -> list:
@@ -1580,7 +1582,7 @@ def __init__(
 
         # Allocate buffers
         if device is None:
-            device = torch.device("cuda")
+            device = torch.device(te_device_type())
 
     def make_quantizers(self) -> list:
         # TODO(ksivamani); Find better design for this, adding here to avoid circular import.
@@ -1638,7 +1640,7 @@ def __init__(
 
         # Allocate buffers
         if device is None:
-            device = torch.device("cuda")
+            device = torch.device(te_device_type())
         self.device = device
 
     def make_quantizers(self) -> list:
@@ -1721,7 +1723,7 @@ def __init__(
 
         # Allocate buffers
         if device is None:
-            device = torch.device("cuda")
+            device = torch.device(te_device_type())
 
     def make_quantizers(self) -> list:
         """Build one ``NVFP4Quantizer`` per slot, dispatched by tensor type.
@@ -1965,7 +1967,7 @@ def __init__(
         self.num_quantizers = num_quantizers
         self.roles = roles
         if device is None:
-            device = torch.device("cuda")
+            device = torch.device(te_device_type())
         self.device = device
 
         # -- Stateful sub-state slots (initialized empty) --
diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py
index f6f2794bdf9..1635cd70682 100644
--- a/transformer_engine/pytorch/quantized_tensor.py
+++ b/transformer_engine/pytorch/quantized_tensor.py
@@ -18,6 +18,8 @@
 
 from transformer_engine.common.recipe import Recipe
 from transformer_engine.pytorch.constants import dist_group_type
+
+from transformer_engine import te_device_type
 from transformer_engine.pytorch.tensor._quantization_helpers import (
     _QuantizeFunc,
     _IdentityFunc,
@@ -501,7 +503,7 @@ def make_empty(
             )
 
         if device is None:
-            device = torch.device("cuda")
+            device = torch.device(te_device_type())
         # Handle the device passed as string
         device = torch.device(device)
         result = tex.create_empty_quantized_tensor(
@@ -557,7 +559,7 @@ def alloc_tensors(
         Returns ``{attr_name: torch.Tensor}`` suitable as the ``inner_tensors``
         argument of the storage's ``__tensor_unflatten__``.
         """
-        device = torch.device(device if device is not None else "cuda")
+        device = torch.device(device if device is not None else te_device_type())
         return {
             attr: torch.empty(buf_shape, dtype=buf_dtype, device=device)
             for attr, (buf_shape, buf_dtype) in self.inner_tensor_specs(tuple(shape)).items()
diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py
index 6105b18a736..c85577fe2f8 100644
--- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py
+++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py
@@ -11,6 +11,8 @@
 
 import torch
 import transformer_engine_torch as tex
+
+from transformer_engine import te_device_type
 from transformer_engine.common.recipe import Float8BlockScaling, Recipe
 from .storage.float8_blockwise_tensor_storage import (
     Float8BlockwiseQTensorStorage,
@@ -441,7 +443,7 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None):
                 qt._rowwise_scale_inv,
                 qt._columnwise_scale_inv,
             ):
-                if t is not None and t.is_cuda:
+                if t is not None and t.device.type == te_device_type():
                     t.record_stream(stream)
             return None
 
@@ -538,7 +540,7 @@ def _set_data(self, tensor: torch.Tensor) -> None:
 
         """
         # Tensor device
-        new_device = tensor.device if tensor.is_cuda else self.device
+        new_device = tensor.device if tensor.device.type == te_device_type() else self.device
 
         def _set_from_tensor(dst: Float8BlockwiseQTensor, src: Float8BlockwiseQTensor):
             dst._rowwise_data = src._rowwise_data
diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py
index cf37c36c598..d37b8e2800e 100644
--- a/transformer_engine/pytorch/tensor/float8_tensor.py
+++ b/transformer_engine/pytorch/tensor/float8_tensor.py
@@ -10,6 +10,7 @@
 from torch.distributed.fsdp._fully_shard._fsdp_common import TrainingState
 import transformer_engine_torch as tex
 
+from transformer_engine import te_device_type
 from transformer_engine.common.recipe import (
     DelayedScaling,
     Float8CurrentScaling,
@@ -1139,7 +1140,7 @@ def _set_data(self, tensor: torch.Tensor) -> None:
         """
 
         # Tensor device
-        new_device = tensor.device if tensor.is_cuda else self.device
+        new_device = tensor.device if tensor.device.type == te_device_type() else self.device
         if not devices_match(new_device, tensor.device):
             tensor = tensor.to(device=new_device)
 
diff --git a/transformer_engine/pytorch/tensor/grouped_tensor.py b/transformer_engine/pytorch/tensor/grouped_tensor.py
index 65b4a848228..474d932165d 100644
--- a/transformer_engine/pytorch/tensor/grouped_tensor.py
+++ b/transformer_engine/pytorch/tensor/grouped_tensor.py
@@ -10,6 +10,7 @@
 import torch
 from torch.utils._pytree import tree_map
 
+from transformer_engine import te_device_type
 from ..quantized_tensor import QuantizedTensorStorage, Quantizer
 from .storage.grouped_tensor_storage import GroupedTensorStorage
 from ..constants import DType
@@ -133,7 +134,7 @@ def __new__(
                 device = maybe_tensor.device
                 break
         if device is None:
-            device = torch.device("cuda")
+            device = torch.device(te_device_type())
 
         # Match QuantizedTensor __new__: accept externally-computed stride to
         # avoid Python-side stride computation overhead for C++ construction.
diff --git a/transformer_engine/pytorch/tensor/identity_tensor.py b/transformer_engine/pytorch/tensor/identity_tensor.py
index ec171564fe2..8f8bda44cbc 100644
--- a/transformer_engine/pytorch/tensor/identity_tensor.py
+++ b/transformer_engine/pytorch/tensor/identity_tensor.py
@@ -19,6 +19,7 @@
 import torch
 from torch.ops import aten
 
+from transformer_engine import te_device_type
 from .storage.identity_tensor_storage import IdentityTensorStorage
 from ..quantized_tensor import QuantizedTensor, QuantizedTensorStorage, Quantizer
 
@@ -114,7 +115,7 @@ def make_empty(
         pin_memory: bool = False,
     ) -> Union["IdentityTensor", IdentityTensorStorage]:
         if device is None:
-            device = torch.device("cuda")
+            device = torch.device(te_device_type())
         device = torch.device(device)
         data_dtype = self.dtype if self.dtype is not None else dtype
         data = torch.empty(tuple(shape), dtype=data_dtype, device=device, pin_memory=pin_memory)
diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py
index 267806e43ee..c4081b75c64 100644
--- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py
+++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py
@@ -13,6 +13,7 @@
 from torch.distributed.fsdp._fully_shard._fsdp_common import TrainingState
 import transformer_engine_torch as tex
 
+from transformer_engine import te_device_type
 from transformer_engine.common.recipe import MXFP8BlockScaling, Recipe
 from ..constants import MXFP8_BLOCK_SCALING_SIZE, DType
 from ..utils import devices_match, round_up_to_nearest_multiple
@@ -823,7 +824,7 @@ def _set_data(self, tensor: torch.Tensor) -> None:
         """
 
         # Tensor device
-        new_device = tensor.device if tensor.is_cuda else self.device
+        new_device = tensor.device if tensor.device.type == te_device_type() else self.device
         if not devices_match(new_device, tensor.device):
             tensor = tensor.to(device=new_device)
 
diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py
index 938d9a77d4c..29be015b6ec 100644
--- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py
+++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py
@@ -13,6 +13,7 @@
 import torch
 import transformer_engine_torch as tex
 
+from transformer_engine import te_device_type
 from transformer_engine.common.recipe import NVFP4BlockScaling, Recipe
 from ..constants import NVFP4_BLOCK_SCALING_SIZE, dist_group_type, DType
 from ..utils import (
@@ -796,7 +797,7 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None):
                 qt._amax_rowwise,
                 qt._amax_columnwise,
             ):
-                if t is not None and t.is_cuda:
+                if t is not None and t.device.type == te_device_type():
                     t.record_stream(stream)
             return None
 
@@ -954,7 +955,7 @@ def _set_data(self, tensor: torch.Tensor) -> None:
         """
 
         # Tensor device
-        new_device = tensor.device if tensor.is_cuda else self.device
+        new_device = tensor.device if tensor.device.type == te_device_type() else self.device
         if not devices_match(new_device, tensor.device):
             tensor = tensor.to(device=new_device)
 
diff --git a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py
index f4a95607d48..cccd048cb83 100644
--- a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py
+++ b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py
@@ -12,6 +12,7 @@
 
 import transformer_engine_torch as tex
 
+from transformer_engine import te_device_type
 from ...quantized_tensor import InnerTensor, QuantizedTensorStorage, Quantizer
 from .._quantization_helpers import safe_quantized_repr
 
@@ -44,8 +45,8 @@ def forward(
         te_dtype = torch_to_transformer_engine_dtype[dtype]
         # ``tex.dequantize`` requires CUDA-resident buffers.
         src_device = tensor.device
-        if src_device.type != "cuda":
-            cuda_tensor = tensor.to(device=torch.device("cuda"))
+        if src_device.type != te_device_type():
+            cuda_tensor = tensor.to(device=torch.device(te_device_type()))
             result = tex.dequantize(cuda_tensor, te_dtype)
             return result.to(device=src_device)
         return tex.dequantize(tensor, te_dtype)
diff --git a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py
index 97b66a8341f..9931325eec3 100644
--- a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py
+++ b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py
@@ -15,6 +15,7 @@
 
 import transformer_engine_torch as tex
 
+from transformer_engine import te_device_type
 from ...quantized_tensor import InnerTensor, QuantizedTensorStorage, Quantizer
 from .._quantization_helpers import safe_quantized_repr
 
@@ -54,8 +55,8 @@ def forward(
 
         # ``tex.dequantize`` requires CUDA-resident buffers. If the tensor has
         src_device = tensor.device
-        if src_device.type != "cuda":
-            cuda_tensor = tensor.to(device=torch.device("cuda"))
+        if src_device.type != te_device_type():
+            cuda_tensor = tensor.to(device=torch.device(te_device_type()))
             result = tex.dequantize(cuda_tensor, torch_to_transformer_engine_dtype[dtype])
             return result.to(device=src_device)
         return tex.dequantize(tensor, torch_to_transformer_engine_dtype[dtype])
diff --git a/transformer_engine/pytorch/transformer.py b/transformer_engine/pytorch/transformer.py
index db9a352881c..cc361454a4c 100644
--- a/transformer_engine/pytorch/transformer.py
+++ b/transformer_engine/pytorch/transformer.py
@@ -14,6 +14,8 @@
 from transformer_engine.pytorch.module import LayerNormMLP, LayerNorm, RMSNorm
 from transformer_engine.pytorch.attention.multi_head_attention import MultiheadAttention
 from transformer_engine.pytorch.attention.inference import InferenceParams
+
+from transformer_engine import te_device_type
 from transformer_engine.pytorch.jit import (
     set_jit_fusion_options,
     warmup_jit_bias_dropout_add_all_dtypes,
@@ -25,6 +27,7 @@
     cast_if_needed,
     get_default_init_method,
     torch_get_autocast_gpu_dtype,
+    torch_is_autocast_gpu_enabled,
 )
 from transformer_engine.pytorch.constants import (
     AttnMaskTypes,
@@ -199,7 +202,7 @@ class TransformerLayer(torch.nn.Module):
                         supports ``'limit'``, ``'alpha'``, and ``'glu_linear_offset'`` parameters.
                         You can set these as
                         ``activation_params={'limit': 7.0, 'alpha': 1.702, 'glu_linear_offset': 1.0}``.
-    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.
@@ -352,7 +355,7 @@ def __init__(
         activation: str = "gelu",
         activation_params: Optional[dict] = None,
         normalization: str = "LayerNorm",
-        device: Union[torch.device, str] = "cuda",
+        device: Union[torch.device, str] = te_device_type(),
         attn_input_format: str = "sbhd",
         name: str = None,
         qk_norm_type: Optional[str] = None,
@@ -907,7 +910,7 @@ def forward(
                 )
 
         # For AMP
-        if torch.is_autocast_enabled():
+        if torch_is_autocast_gpu_enabled():
             hidden_states = cast_if_needed(hidden_states, torch_get_autocast_gpu_dtype())
 
         # Self attention.
diff --git a/transformer_engine/pytorch/triton/mhc.py b/transformer_engine/pytorch/triton/mhc.py
index 41f71966a23..ce9a507b448 100644
--- a/transformer_engine/pytorch/triton/mhc.py
+++ b/transformer_engine/pytorch/triton/mhc.py
@@ -25,6 +25,7 @@
     _mhc_sinkhorn_bwd_fused,
 )
 from transformer_engine.pytorch.cpp_extensions.gemm import general_gemm
+from transformer_engine import te_device_type
 
 ENFORCE_DETERMINISTIC = os.environ.get("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1") == "0"
 
@@ -50,7 +51,7 @@ def _init_tma_allocator():
     def alloc_fn(
         size: int, alignment: int, stream: Optional[int]
     ):  # pylint: disable=unused-argument
-        return torch.empty(size, device="cuda", dtype=torch.int8)
+        return torch.empty(size, device=te_device_type(), dtype=torch.int8)
 
     triton.set_allocator(alloc_fn)
     _tma_allocator_initialized = True
diff --git a/transformer_engine/pytorch/triton/permutation.py b/transformer_engine/pytorch/triton/permutation.py
index c155d73e1e3..744ab751058 100644
--- a/transformer_engine/pytorch/triton/permutation.py
+++ b/transformer_engine/pytorch/triton/permutation.py
@@ -9,6 +9,7 @@
 import torch
 import triton
 
+from transformer_engine import te_device_type
 from transformer_engine.common.triton.permutation import (
     _row_id_map_pass_1_kernel,
     _row_id_map_pass_2_kernel,
@@ -49,10 +50,12 @@ def make_row_id_map(
         The [num_experts, num_experts + n_routed) items are the indices of the experts corresponding
         to the first n_routed row indices above.
     """
-    row_id_map = torch.empty((num_tokens, num_experts * 2 + 1), dtype=torch.int32, device="cuda")
+    row_id_map = torch.empty(
+        (num_tokens, num_experts * 2 + 1), dtype=torch.int32, device=routing_map.device
+    )
     block_size = 1024
     grid = (num_experts, triton.cdiv(num_tokens, block_size))
-    workspace_tensor = torch.empty(grid, dtype=torch.int32, device="cuda")
+    workspace_tensor = torch.empty(grid, dtype=torch.int32, device=routing_map.device)
 
     # supposing num_tokens == 5, num_experts == 3, block_size == 3
     # and we have a routing_map like this:
@@ -160,12 +163,14 @@ def permute_with_mask_map(
     # Use torch.zeros when pad_offsets is provided to ensure padding regions are zeroed.
     # The kernel writes only to valid positions, leaving padding positions at zero.
     alloc = torch.zeros if pad_offsets is not None else torch.empty
-    output = alloc((num_out_tokens, hidden_size), dtype=inp.dtype, device="cuda")
+    output = alloc((num_out_tokens, hidden_size), dtype=inp.dtype, device=te_device_type())
     permuted_probs = (
-        alloc((num_out_tokens,), dtype=probs.dtype, device="cuda") if probs is not None else None
+        alloc((num_out_tokens,), dtype=probs.dtype, device=te_device_type())
+        if probs is not None
+        else None
     )
     permuted_scale = (
-        alloc((num_out_tokens, scale_hidden_dim), dtype=scale.dtype, device="cuda")
+        alloc((num_out_tokens, scale_hidden_dim), dtype=scale.dtype, device=te_device_type())
         if scale is not None
         else None
     )
@@ -243,10 +248,10 @@ def unpermute_with_mask_map(
     hidden_size : int
         Hidden size of the permuted tensor.
     """
-    output = torch.empty((num_tokens, hidden_size), dtype=inp.dtype, device="cuda")
+    output = torch.empty((num_tokens, hidden_size), dtype=inp.dtype, device=te_device_type())
     if permuted_probs is not None:
         unpermuted_probs = torch.empty(
-            (num_tokens, num_experts), dtype=permuted_probs.dtype, device="cuda"
+            (num_tokens, num_experts), dtype=permuted_probs.dtype, device=te_device_type()
         )
     else:
         unpermuted_probs = None
@@ -325,9 +330,11 @@ def unpermute_with_mask_map_bwd_with_merging_probs(
     # by the kernel. This matches the behavior of Fp8Unpadding.backward which zeros
     # out the padding slots.
     alloc = torch.zeros if pad_offsets is not None else torch.empty
-    act_grad = alloc((num_out_tokens, hidden_size), dtype=fwd_output_grad.dtype, device="cuda")
+    act_grad = alloc(
+        (num_out_tokens, hidden_size), dtype=fwd_output_grad.dtype, device=te_device_type()
+    )
     merging_probs_grad = torch.empty(
-        (num_tokens, num_experts), dtype=merging_probs.dtype, device="cuda"
+        (num_tokens, num_experts), dtype=merging_probs.dtype, device=te_device_type()
     )
     grid = (num_tokens,)
     _unpermute_bwd_with_merging_probs_kernel[grid](
@@ -378,7 +385,7 @@ def make_chunk_sort_map(
     num_splits : int
         Number of splits of split_sizes and sorted_indices.
     """
-    row_id_map = torch.empty((num_tokens,), dtype=torch.int32, device="cuda")
+    row_id_map = torch.empty((num_tokens,), dtype=torch.int32, device=te_device_type())
     grid = (num_tokens,)
     _make_chunk_sort_map_kernel[grid](
         split_sizes,
@@ -416,9 +423,9 @@ def sort_chunks_by_map(
     is_forward : bool
         Whether the sort is for forward or backward.
     """
-    output = torch.empty((num_tokens, hidden_size), dtype=inp.dtype, device="cuda")
+    output = torch.empty((num_tokens, hidden_size), dtype=inp.dtype, device=te_device_type())
     if probs is not None:
-        permuted_probs = torch.empty((num_tokens,), dtype=probs.dtype, device="cuda")
+        permuted_probs = torch.empty((num_tokens,), dtype=probs.dtype, device=te_device_type())
     else:
         permuted_probs = None
     # pylint: disable=unnecessary-lambda-assignment
diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py
index 62afd10f153..a3208db5026 100644
--- a/transformer_engine/pytorch/utils.py
+++ b/transformer_engine/pytorch/utils.py
@@ -14,6 +14,7 @@
 import numpy as np
 import torch
 
+from transformer_engine import te_device_type
 from .torch_version import torch_version
 from ..common import decode_cudnn_version
 from ..debug.pytorch.debug_quantization import DebugQuantizedTensor
@@ -789,12 +790,12 @@ def canonicalize_device(device: Optional[torch.device | str]) -> torch.device:
     if device is None:
         # Use default CUDA device
         device = torch.get_default_device()
-        if device.type != "cuda":
-            device = torch.device("cuda", torch.cuda.current_device())
+        if device.type != te_device_type():
+            device = torch.device(te_device_type(), torch.cuda.current_device())
     elif not isinstance(device, torch.device):
         device = torch.device(device)
-    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())
     return device
 
 
@@ -816,7 +817,7 @@ def devices_match(device1: torch.device, device2: torch.device) -> bool:
     device2 = torch.device(device2)
     if device1.type != device2.type:
         return False
-    if device1.type == "cuda":
+    if device1.type == te_device_type():
         index1 = device1.index
         index2 = device2.index
         if index1 == index2:
@@ -952,15 +953,22 @@ def canonicalize_process_group(
     return group
 
 
+def torch_is_autocast_gpu_enabled() -> bool:
+    """Check whether PyTorch autocast is enabled for the GPU device."""
+    if torch_version() >= (2, 4, 0):
+        return torch.is_autocast_enabled(te_device_type())
+    return torch.is_autocast_enabled()
+
+
 def torch_get_autocast_gpu_dtype() -> torch.dtype:
     """Get PyTorch autocast GPU dtype."""
     if torch_version() >= (2, 4, 0):
-        return torch.get_autocast_dtype("cuda")
+        return torch.get_autocast_dtype(te_device_type())
     return torch.get_autocast_gpu_dtype()
 
 
 if torch_version() >= (2, 4, 0):
-    gpu_autocast_ctx = functools.partial(torch.amp.autocast, device_type="cuda")
+    gpu_autocast_ctx = functools.partial(torch.amp.autocast, device_type=te_device_type())
 else:
     gpu_autocast_ctx = torch.cuda.amp.autocast