transformer_engine.h

Transformer Engine API 的基类和函数。

类型定义

typedef void *NVTETensor

TE 张量类型。

NVTETensor 是一种连续的张量类型，存储指向给定形状和类型数据的指针。它不拥有它指向的内存。

枚举

enum NVTEDType

TE 数据类型。

值

enumerator kNVTEByte: 字节

enumerator kNVTEInt32: 32 位整数

enumerator kNVTEInt64: 64 位整数

enumerator kNVTEFloat32: 32 位浮点数

enumerator kNVTEFloat16: 16 位浮点数 (E5M10)

enumerator kNVTEBFloat16: 16 位 bfloat (E8M7)

enumerator kNVTEFloat8E4M3: 8 位浮点数 (E4M3)

enumerator kNVTEFloat8E5M2: 8 位浮点数 (E5M2)

enumerator kNVTEFloat8E8M0: 8 位浮点数 (E8M0)

enumerator kNVTENumTypes: 支持的类型数量

enum NVTETensorParam

指示要设置/获取的张量参数的类型。

值

enumerator kNVTERowwiseData: 可以按行方式使用的数据

enumerator kNVTEColumnwiseData: 可以按列方式使用的数据

enumerator kNVTEScale: 缩放张量

enumerator kNVTEAmax: Amax 张量

enumerator kNVTERowwiseScaleInv: 用于解码行式数据的逆缩放张量

enumerator kNVTEColumnwiseScaleInv: 用于解码列式数据的逆缩放张量

enumerator kNVTENumTensorParams

enum NVTEScalingMode

缩放粒度

值

enumerator NVTE_DELAYED_TENSOR_SCALING: 每个张量一个缩放因子，延迟计算。也用于高精度数据，无需缩放

enumerator NVTE_MXFP8_1D_SCALING: 每 32 个连续元素的块（按行或按列方向）一个缩放因子

enumerator NVTE_INVALID_SCALING

函数

NVTETensor nvte_create_tensor(NVTEScalingMode scaling_mode)

创建一个新的 TE 张量。

创建一个新的 TE 张量。在使用之前，需要设置其参数。TE 张量只是原始数据的包装器，不拥有内存。

参数:: scaling_mode – [in] 张量的缩放模式。
返回值:: 一个新的 TE 张量。

void nvte_destroy_tensor(NVTETensor tensor)

销毁一个 TE 张量。

由于 TE 张量不拥有内存，因此在此操作期间不会释放底层数据。

参数:: tensor – [in] 要销毁的张量。

void *nvte_tensor_data(const NVTETensor tensor)

获取指向张量的行式数据的原始指针。

参数:: tensor – [in] 张量。
返回值:: 指向张量的行式数据的原始指针。

void *nvte_tensor_columnwise_data(const NVTETensor tensor)

获取指向张量的列式数据的原始指针。

参数:: tensor – [in] 张量。
返回值:: 指向张量的列式数据的原始指针。

NVTEShape nvte_tensor_shape(const NVTETensor tensor)

获取张量的数据形状。

参数:: tensor – [in] 张量。
返回值:: 输入张量的形状。

NVTEShape nvte_tensor_columnwise_shape(const NVTETensor tensor)

获取张量的数据形状。

参数:: tensor – [in] 张量。
返回值:: 输入张量的形状。

size_t nvte_tensor_ndims(const NVTETensor tensor)

获取张量的维度数。

参数:: tensor – [in] 张量。
返回值:: 张量维度数。

size_t nvte_tensor_size(const NVTETensor tensor, const size_t dim)

获取特定张量维度的大小。

参数:

tensor – [in] 张量。
size_t – [in] 维度索引。

返回值:

指定维度上张量的大小。

size_t nvte_tensor_numel(const NVTETensor tensor)

获取张量的元素总数。

参数:: tensor – [in] 张量。
返回值:: 张量中的元素数量。

size_t nvte_tensor_element_size(const NVTETensor tensor)

获取张量数据类型的字节大小。

参数:: tensor – [in] 张量。
返回值:: 张量数据类型的字节大小。

NVTEDType nvte_tensor_type(const NVTETensor tensor)

获取张量的数据类型。

参数:: tensor – [in] 张量。
返回值:: 输入张量的数据类型。

float *nvte_tensor_amax(const NVTETensor tensor)

获取指向张量的 amax 数据的指针。

参数:: tensor – [in] 张量。
返回值:: 指向张量的 amax 数据的指针。

float *nvte_tensor_scale(const NVTETensor tensor)

获取指向张量的缩放数据的指针。

参数:: tensor – [in] 张量。
返回值:: 指向张量的缩放数据的指针。

float *nvte_tensor_scale_inv(const NVTETensor tensor)

获取指向张量的逆缩放数据的指针。

参数:: tensor – [in] 张量。
返回值:: 指向张量的逆缩放数据的指针。

NVTEShape nvte_tensor_scale_inv_shape(const NVTETensor tensor)

获取张量的 scale_inv 形状。

参数:: tensor – [in] 张量。
返回值:: 输入张量的 scale_inv 形状。

void nvte_zero_tensor(const NVTETensor tensor, cudaStream_t stream)

将张量值重置为零。

参数:: tensor – [in] 张量。
返回值:: 输入张量的 scale_inv 形状。

void nvte_set_tensor_param(NVTETensor *tensor, NVTETensorParam param_name, const NVTEBasicTensor *param)

设置张量的参数。

参数:

[in/out] – tensor 张量。
param_name – [in] 要设置的参数。
param – [in] 要设置的值。

NVTEBasicTensor nvte_get_tensor_param(const NVTETensor tensor, NVTETensorParam param_name)

获取张量参数的值。

参数:

tensor – [in] 张量。
param_name – [in] 要设置的参数。

NVTEScalingMode nvte_tensor_scaling_mode(const NVTETensor tensor)

获取此张量的缩放粒度。

参数:: tensor – [in] 张量。
返回值:: 包含张量缩放粒度的结构体。

void nvte_tensor_pack_create(NVTETensorPack *pack): 在 NVTETensorPack 中创建 tensors。

void nvte_tensor_pack_destroy(NVTETensorPack *pack): 销毁 NVTETensorPack 中的 tensors。

struct NVTEShape

#include <transformer_engine.h>

张量的形状。

公共成员

const size_t *data: 形状数据，大小为 ndim。

size_t ndim: 维度数。

struct NVTEBasicTensor

#include <transformer_engine.h>

用于填充 NVTETensor 参数的基本张量类型。它不拥有它指向的内存。

公共成员

void *data_ptr

NVTEDType dtype

NVTEShape shape

struct NVTETensorPack

#include <transformer_engine.h>

张量包，通常用于辅助输出。

公共成员

NVTETensor tensors[MAX_SIZE]: 张量的包装器。它们不持有相关的内存。

size_t size = 0: 包中张量的实际数量，0 <= size <= MAX_SIZE。

公共静态属性

static const int MAX_SIZE = 10: 包中张量的最大数量。假定 <= 10。

namespace transformer_engine

包含 Transformer Engine 的 C++ API 的命名空间。

枚举

enum class DType

TE 数据类型。

值

enumerator kByte

enumerator kInt32

enumerator kInt64

enumerator kFloat32

enumerator kFloat16

enumerator kBFloat16

enumerator kFloat8E4M3

enumerator kFloat8E5M2

enumerator kFloat8E8M0

enumerator kNumTypes

struct TensorWrapper

#include <transformer_engine.h>

NVTETensor 类的 C++ 包装器。

公共函数

inline TensorWrapper(void *dptr, const NVTEShape &shape, const DType dtype, float *amax_dptr = nullptr, float *scale_dptr = nullptr, float *scale_inv_dptr = nullptr, const NVTEShape scale_inv_shape = defaultShape, const NVTEScalingMode scaling_mode = NVTE_DELAYED_TENSOR_SCALING)

构造新的 TensorWrapper。

创建具有给定形状、数据类型和数据的新 TE 张量。TE 张量只是原始数据的包装器，不拥有内存。

参数:

dptr – **[输入]** 指向张量数据的指针。
shape – **[输入]** 张量的形状。
dtype – **[输入]** 张量的数据类型。
amax_dptr – **[输入]** 指向 AMAX 值的指针。
scale_dptr – **[输入]** 指向 scale 值的指针。
scale_inv_shape – **[输入]** scale_inv 的形状
scale_inv_dptr – **[输入]** 指向 scale 值的倒数的指针。

inline TensorWrapper(void *dptr, const std::vector<size_t> &shape, const DType dtype, float *amax_dptr = nullptr, float *scale_dptr = nullptr, float *scale_inv_dptr = nullptr, const std::vector<size_t> &scale_inv_shape = {1}, const NVTEScalingMode scaling_mode = NVTE_DELAYED_TENSOR_SCALING)

构造新的 TensorWrapper。

创建具有给定形状、数据类型和数据的新 TE 张量。TE 张量只是原始数据的包装器，不拥有内存。

参数:

dptr – **[输入]** 指向张量数据的指针。
shape – **[输入]** 张量的形状。
dtype – **[输入]** 张量的数据类型。
amax_dptr – **[输入]** 指向 AMAX 值的指针。
scale_dptr – **[输入]** 指向 scale 值的指针。
scale_inv_shape – **[输入]** scale_inv 的形状
scale_inv_dptr – **[输入]** 指向 scale 值的倒数的指针。

inline explicit TensorWrapper(const NVTEScalingMode scaling_mode = NVTE_DELAYED_TENSOR_SCALING)

构造新的空 TensorWrapper。

创建一个新的空 TE 张量，不持有任何内容。

inline ~TensorWrapper(): TensorWrapper 析构函数。

TensorWrapper &operator=(const TensorWrapper &other) = delete

TensorWrapper(const TensorWrapper &other) = delete

inline TensorWrapper(TensorWrapper &&other)

从现有的 TensorWrapper 构造新的 TensorWrapper。

将现有的 TE 张量传递给新的 TensorWrapper。

参数:: other – **[输入/输出]** 数据的来源。

inline TensorWrapper &operator=(TensorWrapper &&other)

从现有的 TensorWrapper 分配数据。

更改现有 TE 张量的所有权。

参数:: other – **[输入/输出]** 数据的来源。

template<typename ShapeType> inline TensorWrapper &set_parameter(const NVTETensorParam param, void *dptr, DType type, const ShapeType &shape) noexcept

template<typename ShapeType> inline TensorWrapper &set_rowwise_data(void *dptr, DType type, const ShapeType &shape) noexcept

template<typename ShapeType> inline TensorWrapper &set_columnwise_data(void *dptr, DType type, const ShapeType &shape) noexcept

template<typename ShapeType> inline TensorWrapper &set_scale(void *dptr, DType type, const ShapeType &shape) noexcept

template<typename ShapeType> inline TensorWrapper &set_amax(void *dptr, DType type, const ShapeType &shape) noexcept

template<typename ShapeType> inline TensorWrapper &set_rowwise_scale_inv(void *dptr, DType type, const ShapeType &shape) noexcept

template<typename ShapeType> inline TensorWrapper &set_columnwise_scale_inv(void *dptr, DType type, const ShapeType &shape) noexcept

inline NVTEBasicTensor get_parameter(const NVTETensorParam param) const noexcept

inline NVTEBasicTensor get_rowwise_data() const noexcept

inline NVTEBasicTensor get_columnwise_data() const noexcept

inline NVTEBasicTensor get_scale() const noexcept

inline NVTEBasicTensor get_amax() const noexcept

inline NVTEBasicTensor get_rowwise_scale_inv() const noexcept

inline NVTEBasicTensor get_columnwise_scale_inv() const noexcept

inline NVTETensor data() const noexcept

获取底层的 NVTETensor。

返回值:: 此 TensorWrapper 持有的 NVTETensor。

inline const NVTEShape shape() const noexcept

获取此 TensorWrapper 的形状。

返回值:: 此 TensorWrapper 的形状。

inline const NVTEShape columnwise_shape() const noexcept

获取此 TensorWrapper 的形状。

返回值:: 此 TensorWrapper 的形状。

inline size_t size(const size_t dim) const

获取此 TensorWrapper 在给定维度上的大小。

参数:: size_t – [in] 维度索引。
返回值:: 此 TensorWrapper 在给定维度上的大小。

inline size_t ndim() const noexcept

获取此 TensorWrapper 的维度数量。

返回值:: 此 TensorWrapper 的维度数量。

inline size_t numel() const noexcept

获取张量中已分配元素的数量。即使 TensorWrapper 具有非零形状，对于具有 nullptr 数据的张量，此函数将返回 0。

返回值:: 张量中的元素数量。

inline size_t element_size() const noexcept

获取张量的元素大小（以字节为单位）。

返回值:: 元素大小（以字节为单位）。

inline size_t bytes() const noexcept

获取张量的已分配大小（以字节为单位）。即使 TensorWrapper 具有非零形状和有效 dtype，对于具有 nullptr 数据的张量，此函数将返回 0。

返回值:: 张量的总大小（以字节为单位）。

inline DType dtype() const noexcept

获取此 TensorWrapper 的数据类型。

返回值:: 此 TensorWrapper 的数据类型。

inline void *dptr() const noexcept

获取指向张量数据的原始指针。

返回值:: 指向张量数据的原始指针。

inline void *columnwise_dptr() const noexcept

获取指向张量数据的原始指针。

返回值:: 指向张量数据的原始指针。

inline float *amax() const noexcept

获取指向张量的 amax 数据的指针。

返回值:: 指向张量的 amax 数据的指针。

inline float *scale() const noexcept

获取指向张量的缩放数据的指针。

返回值:: 指向张量的缩放数据的指针。

inline float *scale_inv() const noexcept

获取指向张量的逆缩放数据的指针。

返回值:: 指向张量的逆缩放数据的指针。

inline const NVTEShape scale_inv_shape() const noexcept

获取此 TensorWrapper 的 scale_inv_shape。

返回值:: 此 TensorWrapper 的 scale_inv_shape。

inline NVTEScalingMode scaling_mode() const noexcept

获取张量的缩放模式。

返回值:: 张量的缩放模式。

inline void zero_(cudaStream_t stream)

公共静态属性

static constexpr size_t defaultData = 1

static constexpr NVTEShape defaultShape = {&defaultData, 1}

Private Functions

inline NVTEShape convertShape(const NVTEShape &s)

inline NVTEShape convertShape(const std::vector<size_t> &s)

Private Members

NVTETensor tensor_ = nullptr: 包装的 NVTETensor。