transformer_engine.h¶

Transformer Engine API 的基类和函数。

类型定义

typedef void *NVTETensor¶

TE 张量类型。

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

枚举

enum NVTEDType¶

TE 数据类型。

值

enumerator kNVTEByte¶: 字节

enumerator kNVTEInt32¶: 32 位整数

enumerator kNVTEInt64¶: 32 位整数

enumerator kNVTEFloat32¶: 32 位浮点数

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

enumerator kNVTEBFloat16¶: 16 位 bfloat (E8M7)

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

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

enumerator kNVTENumTypes¶: 支持的类型数量

函数

NVTETensor nvte_create_tensor(void *dptr, const NVTEShape shape, const NVTEDType dtype, float *amax_dptr, float *scale_dptr, float *scale_inv_dptr)¶

创建一个新的 TE 张量。

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

参数

dptr – [in] 指向张量数据的指针。
shape – [in] 张量的形状。
dtype – [in] 张量的数据类型。
amax_dptr – [in] 指向 AMAX 值的指针。
scale_dptr – [in] 指向比例值的指针。
scale_inv_dptr – [in] 指向比例值的倒数的指针。

返回值

一个新的 TE 张量。

void nvte_destroy_tensor(NVTETensor tensor)¶

销毁一个 TE 张量。

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

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

NVTEDType nvte_tensor_type(const NVTETensor tensor)¶

获取张量的数据类型。

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

NVTEShape nvte_tensor_shape(const NVTETensor tensor)¶

获取张量的数据形状。

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

void *nvte_tensor_data(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] 张量。
返回值: 指向张量的缩放数据倒数的指针。

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 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¶

枚举符 kBFloat16¶

枚举符 kFloat8E4M3¶

枚举符 kFloat8E5M2¶

枚举符 kNumTypes¶

结构体 TensorWrapper¶

#include <transformer_engine.h>

NVTETensor 类的 C++ 封装器。

公共函数

内联 TensorWrapper(void *dptr, const NVTEShape &shape, const DType dtype, float *amax_dptr = nullptr, float *scale_dptr = nullptr, float *scale_inv_dptr = nullptr)¶

构造新的 TensorWrapper。

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

参数

dptr – [in] 指向张量数据的指针。
shape – [in] 张量的形状。
dtype – [in] 张量的数据类型。
amax_dptr – [in] 指向 AMAX 值的指针。
scale_dptr – [in] 指向比例值的指针。
scale_inv_dptr – [in] 指向比例值的倒数的指针。

内联 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)¶

构造新的 TensorWrapper。

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

参数

dptr – [in] 指向张量数据的指针。
shape – [in] 张量的形状。
dtype – [in] 张量的数据类型。
amax_dptr – [in] 指向 AMAX 值的指针。
scale_dptr – [in] 指向比例值的指针。
scale_inv_dptr – [in] 指向比例值的倒数的指针。

内联 TensorWrapper()¶

构造新的空 TensorWrapper。

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

内联 ~TensorWrapper()¶: TensorWrapper 析构函数。

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

TensorWrapper(const TensorWrapper &other) = delete¶

内联 TensorWrapper(TensorWrapper &&other)¶

从现有 TensorWrapper 构造新的 TensorWrapper。

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

参数: other – [inout] 数据的来源。

内联 TensorWrapper &operator=(TensorWrapper &&other)¶

从现有 TensorWrapper 分配数据。

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

参数: other – [inout] 数据的来源。

内联 NVTETensor data() const noexcept¶

获取底层的 NVTETensor。

返回值: 此 TensorWrapper 包含的 NVTETensor。

内联 const NVTEShape shape() const noexcept¶

获取此 TensorWrapper 的形状。

返回值: 此 TensorWrapper 的形状。

内联 DType dtype() const noexcept¶

获取此 TensorWrapper 的数据类型。

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

内联 void *dptr() const noexcept¶

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

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

inline float *amax() const noexcept¶

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

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

inline float *scale() const noexcept¶

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

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

inline float *scale_inv() const noexcept¶

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

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

私有成员

NVTETensor tensor_ = nullptr¶: 封装的 NVTETensor。