[PyTorch] Refactor C++ quantizer infrastructure by timmoon10 · Pull Request #1952 · NVIDIA/TransformerEngine

timmoon10 · 2025-07-15T09:02:18Z

Description

This PR makes three changes to the quantizer infrastructure in the transformer_engine_torch extensions:

Consolidate recipe-specific quantization logic in the Quantizer::quantize function. Previously this was duplicated in functions for quantization, activations, normalization, etc.
Force quantized tensors to match the quantizer's usages in the Quantizer::convert_and_update_tensor function, similar to Make quantize_ respect the usages of the quantizer #1836.
Change Quantizer::create_tensor to always return an uninitialized tensor, removing the need for an unnecessary scale reciprocal. For backward compatibility, some quantizer subclasses provide functions for creating initialized tensors.

Arguments for removing the `rowwise_data` arg from `Quantize::create_tensor`

rowwise_data provides an option to provide an already-initialized data buffer. This was implemented to support some use-cases with attention involving QKV fusing and with the Userbuffers buffer (no longer needed after #1711). However, this design has numerous problems:

There is no option to provide column-wise data.
There is no option to provide scales. For FP8 delayed scaling, this means we need to compute the reciprocal of the quantizer's scale, which adds CPU and GPU overhead that is often unnecessary. For FP8 current scaling, MXFP8, and DSv3, this makes this API unusable.
Scales are not consistent between recipes. FP8 requires a per-tensor scale that is shared between the data and transpose. MXFP8 and DSv3 require separate row-wise and column-wise scales. Hierarchical scaling recipes involve multiple scale tensors for both row-wise and column-wise data. Maintaining a API that can generically cover all cases will be... not fun.
The required buffers are recipe-specific and machine-specific, so if you are using this API you have to already know the recipe. You can just dynamic_cast the quantizer to specific concrete class, so there is not much benefit in a generic API.

This PR removes the rowwise_data arg entirely from the base class, so calling create_tensor will create a tensor with pure uninitialized buffers. NoneQuantizer and Float8Quantizer still expose variants of create_tensor for providing pre-initialized buffers, with better recipe-specific logic.

#1950 is an alternative attempt to avoid the problems of the rowwise_data API in the FP8 current-scaling quantizer. #1836 adds an optional out arg to Quantizer::create_tensor and will force any provided tensor to match the quantizer's usages.

Closes#1836. Closes#1950.

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

Make Quantizer::create_tensor construct uninitialized tensors, with some sub-class variants for constructing initialized tensors
tex.quantize forces quantized tensors to match the quantizer's usages
Consolidate quantization logic in Quantizer::quantize
Support all quantization schemes in activation forward and backward

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

Signed-off-by: zhongboz <zhongboz@nvidia.com>