[AFQ] Optimize tensor_flatten for runtime (#1951)

[ghstack-poisoned]

Author: IvanKobzarev

torchao/dtypes/affine_quantized_tensor.py

@@ -165,14 +165,18 @@ def dequantize(self, output_dtype: Optional[torch.dtype] = None) -> torch.Tensor
             return dq
 
     def __tensor_flatten__(self):
-        return ["tensor_impl"], [
-            self.block_size,
-            self.shape,
-            self.quant_min,
-            self.quant_max,
-            self.zero_point_domain,
-            self.dtype,
-        ]
+        # This is used in rumtime to unwrap AffineQuantizedTensor activations.
+        # AffineQuantizedTensor has __torch_function__ override:
+        # Each getattr will go through it, which is up to 10x slower than default attribute access.
+        with torch._C.DisableTorchFunctionSubclass():
+            return ["tensor_impl"], [
+                self.block_size,
+                self.shape,
+                self.quant_min,
+                self.quant_max,
+                self.zero_point_domain,
+                self.dtype,
+            ]
 
     @classmethod
     def __tensor_unflatten__(