Implementation:NVIDIA TransformerEngine Core Tensor Impl
Appearance
| Field | Value |
|---|---|
| Sources | TransformerEngine |
| Domains | Deep_Learning, Optimization |
| Last Updated | 2026-02-07 14:00 GMT |
Overview
Implements the core C API for tensor management in Transformer Engine, including tensor creation/destruction, shape/scale validation, data type conversions, and the NVTETensor/NVTEGroupedTensor handle system.
Description
transformer_engine.cpp is one of the most critical files in the library, implementing the tensor abstraction layer that all other components depend on:
- Type utilities:
typeToSize/typeToNumBitsfor dtype arithmetic,to_stringfor dtype and scaling mode formatting. - Validation functions:
CheckInputTensor,CheckOutputTensor,CheckScaleTensorShapeverify tensor shapes, scale_inv dimensions, and FP8 block scaling alignment for all scaling modes (tensor, MXFP8, block 1D/2D, NVFP4). - C API implementation:
nvte_create_tensor,nvte_destroy_tensor,nvte_tensor_pack, etc., which allocate/deallocateTensorobjects behind opaqueNVTETensorhandles. - Memory management:
TensorAllocatorandGroupedTensorAllocatorclasses manage allocation with proper alignment.
Usage
This implementation backs all NVTETensor operations. It is called indirectly whenever tensors are created, queried, or validated throughout the library.
Code Reference
Source Location
- Repository
NVIDIA/TransformerEngine- File
transformer_engine/common/transformer_engine.cpp- Lines
- 1--1270
Signature
namespace transformer_engine {
size_t typeToNumBits(const DType type);
size_t typeToSize(const DType type);
std::string to_string(const DType type);
std::string to_string(const NVTEScalingMode &mode);
void CheckInputTensor(const Tensor &t, const std::string &name);
void CheckOutputTensor(Tensor &t, const std::string &name);
void CheckScaleTensorShape(const Tensor &t, const std::string &name);
void CheckNoopTensor(const Tensor &t, const std::string &name);
} // namespace transformer_engine
// C API implementations
NVTETensor nvte_create_tensor(NVTEScalingMode scaling_mode);
void nvte_destroy_tensor(NVTETensor tensor);
Import
#include <transformer_engine/transformer_engine.h>
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
scaling_mode |
NVTEScalingMode |
Yes | Scaling mode for new tensor |
tensor |
NVTETensor |
Yes | Tensor handle for queries/destruction |
Outputs
| Name | Type | Description |
|---|---|---|
NVTETensor |
void* |
Opaque handle to tensor (for creation) |
| type info | various | Type sizes, string representations (for queries) |
Usage Examples
#include <transformer_engine/transformer_engine.h>
// Create a tensor with delayed tensor scaling
NVTETensor tensor = nvte_create_tensor(NVTE_DELAYED_TENSOR_SCALING);
// Query tensor properties
NVTEShape shape = nvte_tensor_shape(tensor);
NVTEDType dtype = nvte_tensor_dtype(tensor);
// Validate tensor configuration
CheckInputTensor(internal_tensor, "input_activation");
// Cleanup
nvte_destroy_tensor(tensor);
Related Pages
Page Connections
Double-click a node to navigate. Hold to expand connections.
Principle
Implementation
Heuristic
Environment