Support quantized tensor subclasses across Python and C++

tests/pytorch/test_hybrid_quantization.py

@@ -1738,6 +1738,25 @@ def test_detach(self, hybrid_tensor):
         assert isinstance(detached, HybridQuantizedTensor)
         assert not detached.requires_grad
 
+    def test_detach_preserves_subclass(self, hybrid_tensor):
+        """HybridQuantizedTensor detach preserves its runtime subclass."""
+
+        class DerivedHybridQuantizedTensor(HybridQuantizedTensor):
+            pass
+
+        hybrid_tensor.__class__ = DerivedHybridQuantizedTensor
+        source_data = hybrid_tensor.get_data_tensors()
+
+        detached = hybrid_tensor.detach()
+
+        assert type(detached) is DerivedHybridQuantizedTensor
+        for detached_data, source_data_tensor in zip(detached.get_data_tensors(), source_data):
+            assert detached_data is source_data_tensor
+        assert not detached.requires_grad
+
+        parameter = torch.nn.Parameter(hybrid_tensor)
+        assert type(parameter) is DerivedHybridQuantizedTensor
+
     def test_repr(self, hybrid_tensor):
         r = repr(hybrid_tensor)
         assert "HybridQuantizedTensor" in r

tests/pytorch/test_identity_quantizer.py

@@ -237,6 +237,26 @@ def test_quantize_returns_identity_tensor(self):
         out = IdentityQuantizer()(x)
         assert isinstance(out, IdentityTensor)
 
+    def test_detach_preserves_subclass(self):
+        """IdentityTensor detach preserves its runtime subclass."""
+
+        class DerivedIdentityTensor(IdentityTensor):
+            pass
+
+        tensor = IdentityQuantizer()(torch.randn(8, 16, device="cuda", dtype=torch.bfloat16))
+        tensor.__class__ = DerivedIdentityTensor
+
+        detached = tensor.detach()
+
+        assert type(detached) is DerivedIdentityTensor
+        assert detached._hp_data.data_ptr() == tensor._hp_data.data_ptr()
+        assert detached._hp_data.stride() == tensor._hp_data.stride()
+        assert detached._hp_data.storage_offset() == tensor._hp_data.storage_offset()
+        assert not detached.requires_grad
+
+        parameter = torch.nn.Parameter(tensor)
+        assert type(parameter) is DerivedIdentityTensor
+
     def test_internal_returns_storage(self):
         x = torch.randn(8, 16, device="cuda", dtype=torch.bfloat16)
         q = IdentityQuantizer()

tests/pytorch/test_quantized_tensor.py

@@ -563,6 +563,86 @@ def setup_class(cls) -> None:
         torch.manual_seed(seed)
         torch.cuda.manual_seed(seed)
 
+    @pytest.mark.parametrize(
+        "quantization",
+        [
+            pytest.param(
+                "fp8",
+                marks=pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8),
+            ),
+            pytest.param(
+                "fp8_blockwise",
+                marks=pytest.mark.skipif(
+                    not fp8_block_scaling_available,
+                    reason=reason_for_no_fp8_block_scaling,
+                ),
+            ),
+            pytest.param(
+                "mxfp8",
+                marks=pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8),
+            ),
+            pytest.param(
+                "nvfp4",
+                marks=pytest.mark.skipif(not nvfp4_available, reason=reason_for_no_nvfp4),
+            ),
+        ],
+    )
+    def test_detach_preserves_subclass(self, quantization: str) -> None:
+        """Detaching a quantized tensor preserves its runtime subclass."""
+        quantizer = make_quantizer(quantization)
+        tensor = quantizer(torch.randn(128, 128, dtype=torch.bfloat16, device="cuda"))
+        derived_type = type(f"Derived{type(tensor).__name__}", (type(tensor),), {})
+        tensor.__class__ = derived_type
+
+        detached = tensor.detach()
+
+        assert type(detached) is derived_type
+        for detached_data, source_data in zip(
+            detached.get_data_tensors(), tensor.get_data_tensors()
+        ):
+            assert detached_data is source_data
+        assert not detached.requires_grad
+
+        parameter = torch.nn.Parameter(tensor)
+        assert type(parameter) is derived_type
+
+    @pytest.mark.parametrize(
+        "quantization",
+        [
+            pytest.param(
+                "fp8",
+                marks=pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8),
+            ),
+            pytest.param(
+                "fp8_blockwise",
+                marks=pytest.mark.skipif(
+                    not fp8_block_scaling_available,
+                    reason=reason_for_no_fp8_block_scaling,
+                ),
+            ),
+            pytest.param(
+                "mxfp8",
+                marks=pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8),
+            ),
+            pytest.param(
+                "nvfp4",
+                marks=pytest.mark.skipif(not nvfp4_available, reason=reason_for_no_nvfp4),
+            ),
+        ],
+    )
+    def test_update_quantized_supports_subclass(self, quantization: str) -> None:
+        """Quantizers update derived quantized tensor wrappers in-place."""
+        quantizer = make_quantizer(quantization)
+        source = torch.randn(128, 128, dtype=torch.bfloat16, device="cuda")
+        tensor = quantizer(source)
+        derived_type = type(f"Derived{type(tensor).__name__}", (type(tensor),), {})
+        tensor.__class__ = derived_type
+
+        quantizer.update_quantized(torch.zeros_like(source), tensor)
+
+        assert type(tensor) is derived_type
+        torch.testing.assert_close(tensor.dequantize(), torch.zeros_like(source))
+
     @pytest.mark.parametrize("op", ("clone", "view", "reshape", "contiguous"))
     @pytest.mark.parametrize("quantization", _quantization_list)
     def test_identity_op(

transformer_engine/pytorch/csrc/pybind.h

@@ -57,13 +57,15 @@ inline bool IsFloat8CurrentScalingQuantizers(PyObject *obj) {
 }
 
 inline bool IsFloat8Tensor(PyObject *obj) {
-  return Py_TYPE(obj) == Float8TensorPythonClass || Py_TYPE(obj) == Float8TensorStoragePythonClass;
+  return PyObject_TypeCheck(obj, Float8TensorPythonClass) ||
+         PyObject_TypeCheck(obj, Float8TensorStoragePythonClass);
 }
 
 inline bool IsMXFP8Quantizers(PyObject *obj) { return Py_TYPE(obj) == MXFP8QuantizerClass; }
 
 inline bool IsMXFP8Tensor(PyObject *obj) {
-  return Py_TYPE(obj) == MXFP8TensorPythonClass || Py_TYPE(obj) == MXFP8TensorStoragePythonClass;
+  return PyObject_TypeCheck(obj, MXFP8TensorPythonClass) ||
+         PyObject_TypeCheck(obj, MXFP8TensorStoragePythonClass);
 }
 
 inline bool IsFloat8BlockwiseQuantizers(PyObject *obj) {
@@ -73,12 +75,13 @@ inline bool IsFloat8BlockwiseQuantizers(PyObject *obj) {
 inline bool IsNVFP4Quantizers(PyObject *obj) { return Py_TYPE(obj) == NVFP4QuantizerClass; }
 
 inline bool IsFloat8BlockwiseQTensor(PyObject *obj) {
-  return Py_TYPE(obj) == Float8BlockwiseQTensorPythonClass ||
-         Py_TYPE(obj) == Float8BlockwiseQTensorStoragePythonClass;
+  return PyObject_TypeCheck(obj, Float8BlockwiseQTensorPythonClass) ||
+         PyObject_TypeCheck(obj, Float8BlockwiseQTensorStoragePythonClass);
 }
 
 inline bool IsNVFP4Tensor(PyObject *obj) {
-  return Py_TYPE(obj) == NVFP4TensorPythonClass || Py_TYPE(obj) == NVFP4TensorStoragePythonClass;
+  return PyObject_TypeCheck(obj, NVFP4TensorPythonClass) ||
+         PyObject_TypeCheck(obj, NVFP4TensorStoragePythonClass);
 }
 
 TensorWrapper NVTETensorFromFloat8Tensor(py::handle tensor, Quantizer *quantizer);

transformer_engine/pytorch/quantized_tensor.py

@@ -797,13 +797,10 @@ def quantize_(self, tensor: torch.Tensor) -> QuantizedTensor:
     def detach(self) -> QuantizedTensor:
         """Create new quantized tensor with same data
 
-        Output tensor must be detached from the current autograd
-        graph.
+        Output tensor must be detached from the current autograd graph.
 
         """
-        raise NotImplementedError(
-            f"{self.__class__.__name__} class does not implement detach function"
-        )
+        return type(self).make_like(self)
 
     def clear(self):
         """Deallocate this tensor's memory. Typically not needed and must be used carefully"""

transformer_engine/pytorch/tensor/float8_blockwise_tensor.py

@@ -360,10 +360,6 @@ def dequantize(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor:
             return _FromFloat8BlockwiseFunc.apply(self, dequant_dtype)
         return _FromFloat8BlockwiseFunc.forward(None, self, dequant_dtype)
 
-    def detach(self) -> Float8BlockwiseQTensor:
-        # pylint: disable=missing-function-docstring
-        return Float8BlockwiseQTensor.make_like(self)
-
     def clone(self) -> Float8BlockwiseQTensor:
         # pylint: disable=missing-function-docstring
         rowwise_data = None

transformer_engine/pytorch/tensor/float8_tensor.py

@@ -519,10 +519,6 @@ def quantize_(
             return self.quantize_(tensor.dequantize(), noop_flag=noop_flag)
         return super().quantize_(tensor, noop_flag=noop_flag)
 
-    def detach(self) -> Float8Tensor:
-        # pylint: disable=missing-function-docstring
-        return Float8Tensor.make_like(self)
-
     def clone(self) -> Float8Tensor:
         # pylint: disable=missing-function-docstring
         # ``_data`` may be None for columnwise-only sub-storages of a

transformer_engine/pytorch/tensor/hybrid_tensor.py

@@ -469,7 +469,7 @@ def detach(self) -> HybridQuantizedTensor:
                     "HybridQuantizedTensor.detach() does not support storage-only "
                     f"columnwise sub-storage {col_cls.__name__}"
                 )
-        return HybridQuantizedTensor(
+        return self.__class__(
             shape=self.shape,
             dtype=self.dtype,
             rowwise_storage=row,

transformer_engine/pytorch/tensor/identity_tensor.py

@@ -266,7 +266,7 @@ def _wrap_data_view(
         self, data: torch.Tensor, *, requires_grad: Optional[bool] = None
     ) -> "IdentityTensor":
         requires_grad = self.requires_grad if requires_grad is None else requires_grad
-        return IdentityTensor(
+        return self.__class__(
             shape=data.shape,
             dtype=self.dtype,
             hp_data=data,

transformer_engine/pytorch/tensor/mxfp8_tensor.py

@@ -321,11 +321,6 @@ def quantize_(
             return self.quantize_(tensor.dequantize())
         return super().quantize_(tensor, noop_flag=noop_flag)
 
-    def detach(self) -> MXFP8Tensor:
-        # pylint: disable=missing-function-docstring
-        # TODO(ksivamani): Fix the detach bug
-        return MXFP8Tensor.make_like(self)
-
     def clone(self) -> MXFP8Tensor:
         # pylint: disable=missing-function-docstring
         # _rowwise_data may be None for columnwise-only sub-storages (hybrid quantization)

transformer_engine/pytorch/tensor/nvfp4_tensor.py

@@ -528,11 +528,6 @@ def quantize_(
         self._get_quantizer().update_quantized(tensor, self, noop_flag=noop_flag)
         return self
 
-    def detach(self) -> NVFP4Tensor:
-        # pylint: disable=missing-function-docstring
-        # TODO(ksivamani): Fix the detach bug
-        return NVFP4Tensor.make_like(self)
-
     def clone(self) -> NVFP4Tensor:
         # pylint: disable=missing-function-docstring
         assert self._rowwise_data is not None