Implementation:NVIDIA TransformerEngine Core Tensor Impl

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/typeToNumBits for dtype arithmetic, to_string for dtype and scaling mode formatting.
• Validation functions: CheckInputTensor, CheckOutputTensor, CheckScaleTensorShape verify 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/deallocate Tensor objects behind opaque NVTETensor handles.
• Memory management: TensorAllocator and GroupedTensorAllocator classes 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

• Environment:NVIDIA_TransformerEngine_CUDA_Toolkit_Requirements