fix(utils): propagate non_blocking in TorchAOBaseTensor._to_copy and _get_to_kwargs

test/test_utils.py

@@ -56,6 +56,39 @@ def __init__(self, data):
         with self.assertRaisesRegex(NotImplementedError, "arg_types"):
             l.weight = torch.nn.Parameter(MyTensor(l.weight))
 
+    def test_get_to_kwargs_non_blocking(self):
+        """Verify _get_to_kwargs parses and returns the non_blocking flag."""
+
+        class MyTensor(TorchAOBaseTensor):
+            tensor_data_names = ["qdata"]
+            tensor_attribute_names = ["attr", "device"]
+
+            def __new__(cls, qdata, attr="attr", device=None):
+                if device is None:
+                    device = qdata.device
+                kwargs = {"device": device, "dtype": qdata.dtype}
+                r = torch.Tensor._make_wrapper_subclass(cls, qdata.shape, **kwargs)
+                r.qdata = qdata
+                r.attr = attr
+                return r
+
+            def __init__(self, qdata, attr="attr", device=None):
+                pass
+
+        t = MyTensor(torch.randn(4, 4))
+
+        # non_blocking=True is preserved
+        kwargs = t._get_to_kwargs(device="cpu", non_blocking=True)
+        self.assertTrue(kwargs["non_blocking"])
+
+        # non_blocking=False (explicit) is preserved
+        kwargs = t._get_to_kwargs(device="cpu", non_blocking=False)
+        self.assertFalse(kwargs["non_blocking"])
+
+        # default (not specified) → False
+        kwargs = t._get_to_kwargs(device="cpu")
+        self.assertFalse(kwargs["non_blocking"])
+
     def _test_default_impls_helper(self, lp_tensor, lp_tensor_for_copy):
         # get `all_tensor_data_names` and `all_tensor_attribute_names`
         all_tensor_data_names = lp_tensor.tensor_data_names.copy()

torchao/utils.py

@@ -592,15 +592,19 @@ def _(func, types, args, kwargs):
         ):
             kwargs = self._get_to_kwargs(*args[1:], **kwargs)
             device = kwargs.pop("device")
+            non_blocking = kwargs.pop("non_blocking", False)
             tensors = [
-                getattr(self, name).to(device) for name in self.tensor_data_names
+                getattr(self, name).to(device, non_blocking=non_blocking)
+                for name in self.tensor_data_names
             ]
             optional_tensors = []
             if hasattr(self, "optional_tensor_data_names"):
                 for tensor_data_name in self.optional_tensor_data_names:
                     maybe_tensor = getattr(self, tensor_data_name)
                     if maybe_tensor is not None:
-                        optional_tensors.append(maybe_tensor.to(device))
+                        optional_tensors.append(
+                            maybe_tensor.to(device, non_blocking=non_blocking)
+                        )
                     else:
                         optional_tensors.append(None)
 
@@ -691,25 +695,21 @@ class MyTensor(torch.Tensor):
 
 
 def _get_to_kwargs(self, *args, **kwargs):
-    """Helper function to get the device and dtype keyword args for `aten._to_copy.default` op
-    only device and dtype are kept
+    """Helper function to get the device, dtype and non_blocking keyword args for `aten._to_copy.default` op
 
-    Returns: {"device": device, "dtype": dtype}
+    Returns: {"device": device, "dtype": dtype, "non_blocking": non_blocking}
     """
     # `torch._C._nn._parse_to` can't handle `layout` argument
     args = tuple(arg for arg in args if not isinstance(arg, torch.layout))
     if "layout" in kwargs:
         kwargs.pop("layout")
-    # ignoring `non_blocking` and `memory_format` args since these are not
-    # very useful for most of the tensor subclasses
-    # if in the future there are use cases that need these, we'd recommend
-    # to override `_get_to_kwargs` and return these args
-    device, dtype, _, _ = torch._C._nn._parse_to(*args, **kwargs)
+    device, dtype, non_blocking, _ = torch._C._nn._parse_to(*args, **kwargs)
     device = self.device if device is None else device
     dtype = self.dtype if dtype is None else dtype
     kwargs = {
         "device": device,
         "dtype": dtype,
+        "non_blocking": non_blocking,
     }
     return kwargs