Implementation:NVIDIA TransformerEngine PyTorch Csrc Common

Field	Value
Sources	TransformerEngine
Domains	Deep_Learning, PyTorch, Quantization
Last Updated	2026-02-07 14:00 GMT

Overview

Core C++ utility implementation bridging PyTorch tensors with Transformer Engine's internal C++ tensor representation (TensorWrapper).

Description

Provides makeTransformerEngineTensor overloads that convert PyTorch tensors (via pybind11 handles or at::Tensor) into TE TensorWrappers, handling shape extraction, dtype mapping, and quantizer integration. Also provides convert_quantizer to dynamically dispatch Python quantizer objects to their C++ counterparts, and helper functions like convertTorchShape, getTensorShape, convert_shape_back_from_fp4, and FP8 type selection via getTransformerEngineFP8Type.

Usage

Foundational bridge layer called by nearly every C++ extension function to convert between PyTorch and TE tensor formats before invoking CUDA kernels.

Code Reference

Source Location

Repository: NVIDIA/TransformerEngine
File: transformer_engine/pytorch/csrc/common.cpp
Lines: 1--327

Signature

namespace transformer_engine::pytorch {

std::vector<size_t> convert_shape_back_from_fp4(const std::vector<size_t>& shape, bool transpose);
std::vector<size_t> getTensorShape(const at::Tensor& t);
std::vector<size_t> convertTorchShape(const std::vector<int64_t>& shape);

TensorWrapper makeTransformerEngineTensor(py::handle tensor);
TensorWrapper makeTransformerEngineTensor(at::Tensor tensor);
TensorWrapper makeTransformerEngineTensor(
    void* data_ptr, const std::vector<size_t>& shape,
    const transformer_engine::DType type);

std::shared_ptr<Quantizer> convert_quantizer(py::handle quantizer);
transformer_engine::DType getTransformerEngineFP8Type(bool is_forward);

}  // namespace transformer_engine::pytorch

Import

#include "common.h"

I/O Contract

Inputs

Name	Type	Required	Description
tensor	`py::handle` or `at::Tensor`	Yes	PyTorch tensor to convert
quantizer	`py::handle`	No	Python quantizer object for FP8 conversion
shape	`std::vector<size_t>`	No	Shape for raw pointer construction

Outputs

Name	Type	Description
tensor_wrapper	`TensorWrapper`	TE-format tensor wrapper ready for CUDA kernel dispatch
quantizer_ptr	`std::shared_ptr<Quantizer>`	C++ quantizer corresponding to the Python quantizer

Usage Examples

#include "common.h"

// Convert a PyTorch tensor to TE format for kernel dispatch
auto te_input = makeTransformerEngineTensor(torch_input);

// Convert a quantizer for FP8 operations
auto cpp_quantizer = convert_quantizer(py_quantizer);

Related Pages

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment