Support quantized tensor subclasses across Python and C++

[dingqingy-nv]

Description

Transformer Engine quantized tensors can be subclassed, but two boundaries assumed an exact concrete wrapper class:

1. Concrete detach() implementations reconstructed the base TE class, dropping the tensor's runtime subclass. PyTorch then rejected the result when rewrapping it as a torch.nn.Parameter, because Parameter requires detach() to preserve the exact runtime type.
2. The PyTorch C++ bindings used exact Py_TYPE comparisons for quantized tensor outputs. As a result, operations such as MXFP8Quantizer.update_quantized(src, dst) rejected a valid MXFP8Tensor subclass even though Python isinstance(dst, MXFP8Tensor) was true.

Megatron Core GTP exposes both cases because it represents native quantized parameters with dynamic subclasses such as GTP_MXFP8Tensor. The newer module application path in #3153 made the latent detach mismatch visible. The downstream integration is NVIDIA/Megatron-LM#6546.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality not to work as expected)
• Infra/Build change
• Code refactoring

Changes

• Define runtime-type preservation as part of the QuantizedTensor.detach() contract.
• Construct detached tensors through the runtime class for Float8Tensor, Float8BlockwiseQTensor, MXFP8Tensor, NVFP4Tensor, IdentityTensor, and HybridQuantizedTensor.
• Use CPython subtype checks for the Float8, blockwise Float8, MXFP8, and NVFP4 tensor/storage wrapper families in the C++ bindings. Quantizer objects retain their existing exact-type checks.
• Add tests that verify exact subclass preservation, aliased quantized storage, detached autograd state, successful torch.nn.Parameter construction, and update_quantized into dynamic subclasses.

Validation

Built and tested on GB300 from TE commit bf64b4e8b2985ce7ff394b7f3cb240e764b24a3a:

• 10/10 quantized subclass regression cases passed, covering detach and C++ update_quantized across Float8, blockwise Float8, MXFP8, NVFP4, Identity, and Hybrid wrappers.
• 3/3 existing module _apply attribute-preservation tests passed.
• 135/135 MXFP8 2D tests passed.
• Transformer Engine Python and C++ lint passed.
• Clean Megatron Core focused suite passed 10/10 with the downstream detach/update workarounds absent.
• Two-node real-data 1D and 2D MXFP8 proxy runs each completed 10 training iterations plus validation with GTP and native MXFP8 parameter gather.
• The same two-node 1D and 2D runs also passed with NVTE_CUTEDSL_FUSED_GROUPED_MLP=1 after installing the official final cuDNN-frontend 1.27.0 package:
  • 1D W&B run
  • 2D W&B run

The cuDNN-frontend package replacement addresses a separate pre-release package mismatch around optional prob_tensor; it is not part of this PR.

Checklist

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Greptile Summary

The PR extends quantized-tensor subclass support across Python detach operations and C++ type detection.

• Centralizes subclass-preserving detach behavior in QuantizedTensor.
• Preserves runtime subclasses for Hybrid and Identity wrappers.
• Accepts derived quantized tensor and storage wrappers at C++ boundaries.
• Adds regression coverage for detach, parameter construction, storage aliasing, and in-place quantization updates.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/quantized_tensor.py	Centralizes detach reconstruction through the tensor's runtime class.
transformer_engine/pytorch/csrc/pybind.h	Replaces exact wrapper-type comparisons with subtype-aware CPython checks.
transformer_engine/pytorch/tensor/hybrid_tensor.py	Reconstructs detached Hybrid tensors using their runtime subclass.
transformer_engine/pytorch/tensor/identity_tensor.py	Preserves the runtime Identity tensor class when wrapping aliased data.
tests/pytorch/test_quantized_tensor.py	Adds coverage for subclass-preserving detach and C++ in-place updates across four quantization formats.

Sequence Diagram

sequenceDiagram
participant Caller
participant Tensor as QuantizedTensor subclass
participant Dispatch as PyTorch dispatch
participant Binding as C++ binding
Caller->>Dispatch: detach(tensor)
Dispatch->>Tensor: tensor.detach()
Tensor->>Tensor: type(self).make_like(self)
Tensor-->>Caller: detached runtime subclass
Caller->>Binding: update_quantized(src, subclass)
Binding->>Binding: PyObject_TypeCheck
Binding-->>Caller: update accepted in place

Loading

_{Reviews (2): Last reviewed commit: "Centralize quantized tensor detach" | Re-trigger Greptile}

[zhongbozhu]

LGTM

[ptrendx]

/te-ci pytorch