Implementation:NVIDIA TransformerEngine PyTorch Csrc Common Header
| Field | Value |
|---|---|
| Sources | TransformerEngine |
| Domains | Deep_Learning, PyTorch, Quantization |
| Last Updated | 2026-02-07 14:00 GMT |
Overview
Central C++ header declaring the shared types, classes, and utility functions used across all PyTorch C++ extension files, including the quantizer class hierarchy.
Description
Defines FP8TensorMeta (scale/scale_inv/amax_history storage), the Quantizer abstract base class with its virtual interface (get_scaling_mode, create_tensor, quantize), and concrete quantizer subclasses for every supported format: NoneQuantizer, Float8Quantizer (delayed scaling), Float8CurrentScalingQuantizer, Float8BlockQuantizer, MXFP8Quantizer, and NVFP4Quantizer. Also declares enums FP8FwdTensors/FP8BwdTensors for named indices into FP8 metadata arrays, and utility function signatures for tensor conversion and shape manipulation.
Usage
Core header included by every csrc file. Establishes the quantization abstraction hierarchy and the PyTorch-to-TE type bridge that the entire extension layer depends on.
Code Reference
Source Location
- Repository
NVIDIA/TransformerEngine- File
transformer_engine/pytorch/csrc/common.h- Lines
- 1--541
Signature
namespace transformer_engine::pytorch {
struct FP8TensorMeta {
at::Tensor scale;
at::Tensor scale_inv;
at::Tensor amax_history;
};
class Quantizer {
public:
virtual ~Quantizer() = default;
virtual transformer_engine::ScalingMode get_scaling_mode() const = 0;
virtual TensorWrapper create_tensor(py::handle quantized_tensor, ...) = 0;
virtual py::object quantize(TensorWrapper input, ...) = 0;
};
class NoneQuantizer : public Quantizer { ... };
class Float8Quantizer : public Quantizer { ... };
class Float8CurrentScalingQuantizer : public Quantizer { ... };
class Float8BlockQuantizer : public Quantizer { ... };
class MXFP8Quantizer : public Quantizer { ... };
class NVFP4Quantizer : public Quantizer { ... };
enum FP8FwdTensors { ... };
enum FP8BwdTensors { ... };
} // namespace transformer_engine::pytorch
Import
#include "common.h"
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| N/A | N/A | N/A | Header file -- no runtime inputs; defines types and interfaces |
Outputs
| Name | Type | Description |
|---|---|---|
| N/A | N/A | Header file -- provides type declarations for the C++ extension layer |
Usage Examples
#include "common.h"
using namespace transformer_engine::pytorch;
// Use quantizer to create an FP8 tensor
auto quantizer = std::make_shared<Float8Quantizer>(...);
auto te_tensor = quantizer->create_tensor(py_tensor, ...);