@@ -13,6 +13,7 @@
 import torch
 
 import transformer_engine.pytorch as te
+import transformer_engine.pytorch.ops as te_ops
 from transformer_engine.common.recipe import DelayedScaling, Format
 from transformer_engine.pytorch.quantization import FP8GlobalStateManager
 
@@ -109,6 +110,43 @@ def test_module_off_current_device():
     _assert_state_evolved(module)
 
 
+def test_prepare_forward_exception_restores_current_device(monkeypatch):
+    """A failed prepare must not leak the temporary CUDA device guard."""
+    assert torch.cuda.current_device() == 0
+    module = te.Linear(16, 16, bias=False, params_dtype=DT, device="cuda:1")
+    inp = torch.randn(2, 16, device="cuda:1", dtype=DT)
+
+    def fail_init(*args, **kwargs):
+        raise RuntimeError("injected prepare_forward failure")
+
+    monkeypatch.setattr(module, "init_fp8_metadata", fail_init)
+    with pytest.raises(RuntimeError, match="injected prepare_forward failure"):
+        module.prepare_forward(inp)
+
+    assert torch.cuda.current_device() == 0
+    assert module._forward_device_guards == []
+
+
+def test_basic_operation_device_awareness():
+    """The standalone ops API must use the op's device for state and execution."""
+    recipe = DelayedScaling(fp8_format=Format.HYBRID)
+    assert torch.cuda.current_device() == 0
+    with te.quantized_model_init(enabled=True, recipe=recipe):
+        op = te_ops.basic.BasicLinear(512, 1024, device="cuda:1", dtype=DT)
+    inp = torch.randn(128, 512, device="cuda:1", dtype=DT, requires_grad=True)
+    with te.autocast(enabled=True, recipe=DelayedScaling(fp8_format=Format.HYBRID)):
+        out = op(inp)
+    out.sum().backward()
+    torch.cuda.synchronize()
+    assert out.device == torch.device("cuda:1")
+    assert inp.grad is not None and inp.grad.device == torch.device("cuda:1")
+    assert torch.cuda.current_device() == 0
+    for mode in ("forward", "backward"):
+        state = op._fp8_metas[mode][FP8GlobalStateManager.get_meta_tensor_key(mode == "forward")]
+        assert state.scale.device == torch.device("cuda:1")
+        assert state.amax_history.device == torch.device("cuda:1")
+
+
 def test_two_modules_on_different_devices_one_autocast():
     """Two modules on different devices inside one autocast (#3124 case C)."""
     torch.manual_seed(1234)
 
@@ -1105,9 +1105,7 @@ def adjust_amax_history_length(self, length: int, fwd: Optional[bool] = None) ->
                         qstate.global_amax_buffer[buffer_key][pos] = self.fp8_meta[
                             meta_key
                         ].amax_history[0]
-                        qstate.global_amax_devices[buffer_key].add(
-                            self.fp8_meta[meta_key].amax_history.device
-                        )
+                        FP8GlobalStateManager._refresh_global_amax_devices(buffer_key)
 
     def _recipe_state_device(self) -> torch.device:
         """Device the module's quantization state must be allocated on.
@@ -1127,6 +1125,14 @@ def _recipe_state_device(self) -> torch.device:
                 return tensor.device
         return torch.device("cuda", torch.cuda.current_device())
 
+    def _release_forward_device_guard(self) -> None:
+        """Release the most recently entered forward device guard, if any."""
+        if not self._forward_device_guards:
+            return
+        guard = self._forward_device_guards.pop()
+        if guard is not None:
+            guard.__exit__(None, None, None)
+
     def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None:
         """Init scales and amaxes for fwd | bwd."""
         fp8_meta_tensor_key = "scaling_fwd" if fwd else "scaling_bwd"
@@ -1632,68 +1638,81 @@ def prepare_forward(
             guard = torch.cuda.device(inp.device)
             guard.__enter__()
         self._forward_device_guards.append(guard)
+        nvtx_pushed = False
 
-        self.fast_setattr(
-            "allow_different_data_and_param_types", allow_different_data_and_param_types
-        )
-        self.fast_setattr("forwarded_at_least_once", True)
-
-        # Activation recomputation is used and this is the second forward phase.
-        if self.fp8 and in_fp8_activation_recompute_phase():
-            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:
-                raise RuntimeError(
-                    f"TransformerEngine needs CUDA. Got input on device: {inp.device}"
-                )
+        try:
+            self.fast_setattr(
+                "allow_different_data_and_param_types", allow_different_data_and_param_types
+            )
+            self.fast_setattr("forwarded_at_least_once", True)
 
-            if self.tp_size > 1:
-                if not self.tp_group_initialized:
+            # Activation recomputation is used and this is the second forward phase.
+            if self.fp8 and in_fp8_activation_recompute_phase():
+                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:
                     raise RuntimeError(
-                        "Tensor parallel group not initialized. Call "
-                        "set_tensor_parallel_group() before forward pass when tp_size > 1."
+                        f"TransformerEngine needs CUDA. Got input on device: {inp.device}"
                     )
 
-            self.set_activation_dtype(inp)
-            self.init_fp8_metadata(num_gemms=num_gemms)
-            self._check_weight_tensor_recipe_correspondence()
-
-            delayed_scaling_recipe = self.fp8 and _has_delayed_scaling_state(self.fp8_meta)
-            if delayed_scaling_recipe:
-                if self.sequence_parallel:
-                    assert (
-                        self.fp8_meta["recipe"].custom() or self.fp8_meta["recipe"].reduce_amax
-                    ), (
-                        "Amax reduction across tensor parallel group is "
-                        "necessary when using sequence parallelism with FP8."
-                    )
+                if self.tp_size > 1:
+                    if not self.tp_group_initialized:
+                        raise RuntimeError(
+                            "Tensor parallel group not initialized. Call "
+                            "set_tensor_parallel_group() before forward pass when tp_size > 1."
+                        )
+
+                self.set_activation_dtype(inp)
+                self.init_fp8_metadata(num_gemms=num_gemms)
+                self._check_weight_tensor_recipe_correspondence()
+
+                delayed_scaling_recipe = self.fp8 and _has_delayed_scaling_state(self.fp8_meta)
+                if delayed_scaling_recipe:
+                    if self.sequence_parallel:
+                        assert (
+                            self.fp8_meta["recipe"].custom() or self.fp8_meta["recipe"].reduce_amax
+                        ), (
+                            "Amax reduction across tensor parallel group is "
+                            "necessary when using sequence parallelism with FP8."
+                        )
 
-                if not FP8GlobalStateManager.fp8_graph_capturing():
-                    FP8GlobalStateManager.add_fp8_tensors_to_global_buffer(self.fp8_meta)
+                    if not FP8GlobalStateManager.fp8_graph_capturing():
+                        FP8GlobalStateManager.add_fp8_tensors_to_global_buffer(self.fp8_meta)
 
-                # Activation recomputation is used and this is the first forward phase.
-                if self.training and is_fp8_activation_recompute_enabled():
-                    FP8GlobalStateManager.copy_forward_fp8_meta_tensors_for_recompute(self.fp8_meta)
+                    # Activation recomputation is used and this is the first forward phase.
+                    if self.training and is_fp8_activation_recompute_enabled():
+                        FP8GlobalStateManager.copy_forward_fp8_meta_tensors_for_recompute(
+                            self.fp8_meta
+                        )
 
-        nvtx_range_push(self.__class__.__name__ + " forward")
-        if not allow_non_contiguous and not inp.is_contiguous():
-            inp = inp.contiguous()
-        return inp
+            nvtx_range_push(self.__class__.__name__ + " forward")
+            nvtx_pushed = True
+            if not allow_non_contiguous and not inp.is_contiguous():
+                inp = inp.contiguous()
+            return inp
+        except BaseException:
+            try:
+                if nvtx_pushed:
+                    nvtx_range_pop()
+            finally:
+                self._release_forward_device_guard()
+            raise
 
     def end_forward(self):
         """
         Required to be called at the end of the forward function to properly handle
         DelayedScaling metadata handling and the NVTX ranges.
         """
-        delayed_scaling_recipe = self.fp8 and _has_delayed_scaling_state(self.fp8_meta)
-        if delayed_scaling_recipe and self.fp8 and in_fp8_activation_recompute_phase():
-            FP8GlobalStateManager.restore_fp8_meta_tensors(self.fp8_meta)
-        nvtx_range_pop()
-        if self._forward_device_guards:
-            guard = self._forward_device_guards.pop()
-            if guard is not None:
-                guard.__exit__(None, None, None)
+        try:
+            try:
+                delayed_scaling_recipe = self.fp8 and _has_delayed_scaling_state(self.fp8_meta)
+                if delayed_scaling_recipe and self.fp8 and in_fp8_activation_recompute_phase():
+                    FP8GlobalStateManager.restore_fp8_meta_tensors(self.fp8_meta)
+            finally:
+                nvtx_range_pop()
+        finally:
+            self._release_forward_device_guard()
 
     @contextmanager
     def prepare_forward_ctx(
 
@@ -146,7 +146,10 @@ def __init__(
         # Initialize recipe state if needed for natively quantized weight
         self._with_quantized_weight: bool = FP8GlobalStateManager.with_fp8_parameters()
         if self._with_quantized_weight:
-            self.reset_recipe_state(recipe=FP8GlobalStateManager.get_fp8_recipe())
+            self.reset_recipe_state(
+                recipe=FP8GlobalStateManager.get_fp8_recipe(),
+                device=device if device.type == "cuda" else None,
+            )
 
         # Initialize parameters if needed
         weight = torch.empty(
@@ -333,7 +336,10 @@ def reset_parameters(self) -> None:
                 columnwise=torch.is_grad_enabled(),
             )
             quantizer.internal = False
-            with torch.no_grad():
+            device_context = (
+                torch.cuda.device(weight.device) if weight.is_cuda else contextlib.nullcontext()
+            )
+            with device_context, torch.no_grad():
                 weight = quantizer(weight)
 
         # Save updated parameter
@@ -361,8 +367,10 @@ def pre_fuser_forward(self, *, requires_grad: bool) -> None:
             weight_quantizer.set_usage(rowwise=True, columnwise=requires_grad)
             grad_output_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage)
 
-    def reset_recipe_state(self, *, recipe: Optional[Recipe]) -> None:
-        super().reset_recipe_state(recipe=recipe)
+    def reset_recipe_state(
+        self, *, recipe: Optional[Recipe], device: Optional[torch.device] = None
+    ) -> None:
+        super().reset_recipe_state(recipe=recipe, device=device)
 
         # Configure input/grad output quantizers
         # Note: These tensors are only used internally. If there is no
 
@@ -242,7 +242,10 @@ def __init__(
         # Initialize recipe state if needed for natively quantized weight
         self._with_quantized_weight: bool = FP8GlobalStateManager.with_fp8_parameters()
         if self._with_quantized_weight:
-            self.reset_recipe_state(recipe=FP8GlobalStateManager.get_fp8_recipe())
+            self.reset_recipe_state(
+                recipe=FP8GlobalStateManager.get_fp8_recipe(),
+                device=device if device.type == "cuda" else None,
+            )
 
         # RNG state tracker
         self._rng_state_tracker_function: Optional[Callable[[], CudaRNGStatesTracker]]
@@ -585,7 +588,10 @@ def _quantize_weights(
             return self._quantize_weights_mxfp8(weights, quantizers)
 
         # Use quantizers to construct quantized weights
-        with torch.no_grad():
+        device_context = (
+            torch.cuda.device(weights[0].device) if weights[0].is_cuda else contextlib.nullcontext()
+        )
+        with device_context, torch.no_grad():
             return [quantizer(weight) for quantizer, weight in zip(quantizers, weights)]
 
     def _quantize_weights_mxfp8(
@@ -775,8 +781,10 @@ def pre_fuser_forward(self, *, requires_grad: bool) -> None:
                 weight_quantizer.set_usage(rowwise=True, columnwise=requires_grad)
                 grad_output_quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad)
 
-    def reset_recipe_state(self, *, recipe: Optional[Recipe]) -> None:
-        super().reset_recipe_state(recipe=recipe)
+    def reset_recipe_state(
+        self, *, recipe: Optional[Recipe], device: Optional[torch.device] = None
+    ) -> None:
+        super().reset_recipe_state(recipe=recipe, device=device)
 
         for group_idx in range(self.num_groups):
             # Input/grad output quantizers use internal tensors
 
@@ -266,6 +266,25 @@ def backward(
         func_ctx: Any,
         grad_output: torch.Tensor,
         *grad_extra_outputs: torch.Tensor,
+    ) -> tuple[Optional[torch.Tensor], ...]:
+        """Run backward under the output tensor's CUDA device context."""
+        guard = None
+        if grad_output.is_cuda and grad_output.device.index != torch.cuda.current_device():
+            guard = torch.cuda.device(grad_output.device)
+            guard.__enter__()
+        try:
+            return _OperationFuserAutogradFunction._backward_impl(
+                func_ctx, grad_output, *grad_extra_outputs
+            )
+        finally:
+            if guard is not None:
+                guard.__exit__(None, None, None)
+
+    @staticmethod
+    def _backward_impl(
+        func_ctx: Any,
+        grad_output: torch.Tensor,
+        *grad_extra_outputs: torch.Tensor,
     ) -> tuple[Optional[torch.Tensor], ...]:
         """Backward pass"""
 
@@ -697,6 +716,23 @@ def __call__(
         *extra_inputs: torch.Tensor,
         basic_op_kwargs: Optional[list[dict[str, Any]]] = None,
     ) -> torch.Tensor | tuple[torch.Tensor, ...]:
+        """Apply this fuser under the input tensor's CUDA device context."""
+        guard = None
+        if input.is_cuda and input.device.index != torch.cuda.current_device():
+            guard = torch.cuda.device(input.device)
+            guard.__enter__()
+        try:
+            return self._call_impl(input, *extra_inputs, basic_op_kwargs=basic_op_kwargs)
+        finally:
+            if guard is not None:
+                guard.__exit__(None, None, None)
+
+    def _call_impl(
+        self,
+        input: torch.Tensor,  # pylint: disable=redefined-builtin
+        *extra_inputs: torch.Tensor,
+        basic_op_kwargs: Optional[list[dict[str, Any]]] = None,
+    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
 
         # Verify extra input count
         if len(extra_inputs) != self.num_extra_inputs:
 
@@ -315,10 +315,21 @@ def get_grad_output_quantizer(self) -> Optional[Quantizer]:
             return self.get_quantizer("backward", 0)
         return None
 
+    def _recipe_state_device(self) -> torch.device:
+        """Device on which this operation's quantization state must be allocated."""
+        for tensor in self.parameters(recurse=False):
+            if tensor is not None and tensor.device.type == "cuda":
+                return tensor.device
+        for tensor in self.buffers(recurse=False):
+            if tensor is not None and tensor.device.type == "cuda":
+                return tensor.device
+        return torch.device("cuda", torch.cuda.current_device())
+
     def reset_recipe_state(
         self,
         *,
         recipe: Optional[Recipe],
+        device: Optional[torch.device] = None,
     ) -> None:
         """Construct state for quantization recipe"""
 
@@ -369,6 +380,7 @@ def reset_recipe_state(
                     recipe,
                     mode=mode,
                     num_quantizers=num_quantizers,
+                    device=device if device is not None else self._recipe_state_device(),
                     roles=roles,
                 )
                 fp8_meta_key = FP8GlobalStateManager.get_meta_tensor_key(
@@ -431,9 +443,7 @@ def reset_recipe_state(
                             qstate.global_amax_history_buffer[buffer_key][
                                 pos
                             ] = recipe_state.amax_history
-                            qstate.global_amax_devices[buffer_key].add(
-                                recipe_state.amax_history.device
-                            )
+                            FP8GlobalStateManager._refresh_global_amax_devices(buffer_key)
 
         # Add meta tensors to global buffer to participate in reduction
         for mode in ("forward", "backward"):
 
@@ -573,6 +573,15 @@ def add_fp8_tensors_to_global_buffer(
             fp8_meta[index_in_buffer].append(len(qstate.global_amax_buffer[key]) - 1)
             fp8_meta[index_in_buffer].append(key)
 
+    @classmethod
+    def _refresh_global_amax_devices(cls, buffer_key: str) -> None:
+        """Rebuild the device index after registered amax tensors are replaced."""
+        qstate = cls.quantization_state
+        if buffer_key in qstate.global_amax_buffer:
+            qstate.global_amax_devices[buffer_key] = {
+                tensor.device for tensor in qstate.global_amax_buffer[buffer_key]
+            }
+
     @classmethod
     def is_fp8_enabled(cls) -> bool:
         """Is FP8 enabled"""