[PyTorch] Add dtype information to QuantizedTensorStorage class

[ptrendx]

Description

This PR adds the fake dtype information to the QuantizedTensorStorage class. This eliminates the need to guess the correct type for dequantize, as was the case in the distributed.py, and it eliminates the unintentional dequantization to FP32 when calling dequantize() on the Storage class with no dtype argument.

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 to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Please list the changes introduced in this PR:

• Added the _dtype field to the QuantizedTensorStorage class
• Modified the dequantize call to use that new field when calling dequantize with no arguments
• Removed guessing of the dtype from distributed.py

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 Overview

Greptile Summary

This PR adds dtype tracking to quantized tensor storage classes to eliminate dtype guessing during dequantization. Previously, the system had to guess the target dtype (hardcoded to torch.bfloat16 in distributed.py) or default to torch.float32 when calling dequantize() without arguments.

Key changes:

• Added _dtype field to QuantizedTensorStorage base class to track the "fake" high-precision dtype
• Modified all storage class constructors (Float8TensorStorage, NVFP8TensorStorage, MXFP8TensorStorage, Float8BlockwiseQTensorStorage) to accept and store fake_dtype parameter
• Updated dequantize() methods to use self._dtype as default instead of torch.float32
• Removed dtype guessing in distributed.py (replaced hardcoded torch.bfloat16 with inp._dtype)
• Updated C++ quantizer code to pass fake_dtype using GetATenDType(dtype)
• Ensured fake_dtype is propagated through all tensor operations (view, copy, split)

Implementation approach:
The change maintains backward compatibility by defaulting fake_dtype to torch.float32 when not provided. The QuantizedTensor.__new__() validates that fake_dtype matches dtype to prevent inconsistencies. All storage classes follow a consistent pattern where _dtype is set in __new__() and propagated through operations like view().

This is a clean refactoring that improves type safety and eliminates magic constants throughout the codebase.

Confidence Score: 5/5

• This PR is safe to merge with minimal risk
• The implementation is well-structured with consistent patterns across all storage classes, proper validation, backward compatibility via sensible defaults, and comprehensive coverage of all tensor creation paths including C++ bindings
• No files require special attention

Important Files Changed

Filename	Overview
transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py	Adds fake_dtype parameter to storage class, defaulting to torch.float32, and uses it in dequantize() when dtype not specified
transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py	Adds fake_dtype parameter with same pattern as Float8TensorStorage, properly propagated in view() method
transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py	Adds fake_dtype parameter with consistent implementation across storage classes
transformer_engine/pytorch/distributed.py	Removes hardcoded torch.bfloat16 dtype guessing, now uses inp._dtype from storage classes
transformer_engine/pytorch/quantized_tensor.py	Adds _dtype field to base class, validates fake_dtype matches dtype in __new__, updates dequantize() calls to use default dtype
transformer_engine/pytorch/csrc/quantizer.cpp	Adds fake_dtype parameter to all C++ quantizer create_tensor calls using GetATenDType(dtype)

Sequence Diagram

sequenceDiagram
    participant User
    participant Quantizer
    participant TensorStorage
    participant QuantizedTensor
    participant Dequantize

    User->>Quantizer: quantize(tensor, fake_dtype=bfloat16)
    Quantizer->>TensorStorage: __new__(data, fake_dtype=bfloat16)
    TensorStorage->>TensorStorage: _dtype = fake_dtype or float32
    TensorStorage->>QuantizedTensor: super().__new__(dtype=fake_dtype)
    QuantizedTensor->>QuantizedTensor: validate fake_dtype == dtype
    QuantizedTensor->>QuantizedTensor: _dtype = dtype
    QuantizedTensor-->>User: return quantized_tensor

    User->>QuantizedTensor: dequantize()
    QuantizedTensor->>Dequantize: dequantize(dtype=self._dtype)
    Dequantize-->>User: return high_precision_tensor

Loading

[greptile-apps]

_{13 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[ptrendx]

/te-ci pytorch

[ksivaman]

LGTM

[timmoon10]

Overall this is a big improvement. I have some naming nits.

[timmoon10]

Isn't this redundant with the dtype kwarg?

[timmoon10]

I'd prefer to just name it dtype since QuantizedTensor is already using that name in its constructor.

Suggested change

	fake_dtype: Optional[torch.dtype] = None,
	dtype: Optional[torch.dtype] = None,

[timmoon10]

NVFP4 compatibility with FP32 is not very good, so maybe it would be worth changing the default to BF16? The downside is that this default is different from standard PyTorch and from the other quantized tensors.

Suggested change

	instance._dtype = fake_dtype if fake_dtype is not None else torch.float32
	instance._dtype = fake_dtype if fake_dtype is not None else torch.bfloat16