transpose.h¶
处理转置的函数。
函数
-
void nvte_cast_transpose(const NVTETensor input, NVTETensor cast_output, NVTETensor transposed_output, cudaStream_t stream)¶
转换并转置输入。
此函数转换输入并生成 2 个结果
cast_output是转换的结果transposed_output是转换后的结果。
- 参数
input – [in] 形状为 [N, H] 的输入张量。
cast_output – [inout] 转换的结果。形状：[N, H]。
transposed_output – [inout] 转换和转置的结果。形状：[H, N]。
stream – [in] 用于操作的 CUDA 流。
-
void nvte_transpose(const NVTETensor input, NVTETensor transposed_output, cudaStream_t stream)¶
转置输入。
- 参数
input – [in] 形状为 [N, H] 的输入张量。
transposed_output – [out] 转置的结果。形状：[H, N]。
stream – [in] 用于操作的 CUDA 流。
-
void nvte_cast_transpose_dbias(const NVTETensor input, NVTETensor cast_output, NVTETensor transposed_output, NVTETensor dbias, NVTETensor workspace, cudaStream_t stream)¶
转换并转置输入。 此外，沿第一维度减少输入。
此函数转换输入并产生 3 个结果
cast_output是转换的结果transposed_output是转换后的结果。dbias是输入沿第一维减少的结果。
使用空张量作为 workspace 调用此函数不会执行操作，而是将 workspace 张量的形状和类型设置为所需的值。
- 参数
input – [in] 形状为 [N, H] 的输入张量。
cast_output – [inout] 转换的结果。形状：[N, H]。
transposed_output – [inout] 转换和转置的结果。形状：[H, N]。
dbias – [out] 输入沿第一维减少的结果。形状：[H]。
workspace – [out] 工作区张量。
stream – [in] 用于操作的 CUDA 流。
-
void nvte_fp8_transpose_dbias(const NVTETensor input, NVTETensor transposed_output, NVTETensor dbias, NVTETensor workspace, cudaStream_t stream)¶
转置 FP8 输入。此外，沿第一维度缩减输入。
此函数接受 FP8 输入并产生 2 个结果
transposed_output是输入的转置结果。dbias是输入沿第一维减少的结果。
使用空张量作为 workspace 调用此函数不会执行操作，而是将 workspace 张量的形状和类型设置为所需的值。
- 参数
input – [in] 形状为 [N, H] 的输入张量。
transposed_output – [inout] 转置的结果。形状：[H, N]。
dbias – [out] 输入沿第一维减少的结果。形状：[H]。
workspace – [out] 工作区张量。
stream – [in] 用于操作的 CUDA 流。
-
void nvte_cast_transpose_dbias_dgelu(const NVTETensor input, const NVTETensor gelu_input, NVTETensor cast_output, NVTETensor transposed_output, NVTETensor dbias, NVTETensor workspace, cudaStream_t stream)¶
计算输入上 GELU 操作的反向传播，然后进行类型转换和转置。 此外，沿第一维度缩减 GELU 反向传播的结果。
此函数产生 3 个结果
cast_output等于cast(dGELU(input))transposed_output等于transpose(cast(dGELU(input)))dbias等于reduce(dGELU(input), axis=0)
使用空张量作为 workspace 调用此函数不会执行操作，而是将 workspace 张量的形状和类型设置为所需的值。
- 参数
input – [in] 形状为 [N, H] 的输入张量。
gelu_input – [in] 用作 GELU 操作正向传播输入的张量。形状 [N, H]。
cast_output – [inout] 转换的结果。形状：[N, H]。
transposed_output – [inout] 转换和转置的结果。形状：[H, N]。
dbias – [out] 沿第一维度缩减 dGELU(input) 的结果。形状：[H]。
workspace – [out] 工作区张量。
stream – [in] 用于操作的 CUDA 流。
-
void nvte_multi_cast_transpose(size_t num_tensors, const NVTETensor *input_list, NVTETensor *cast_output_list, NVTETensor *transposed_output_list, cudaStream_t stream)¶
类型转换和转置多个张量。
此函数类型转换每个输入张量并产生 2 个结果
cast_output是转换的结果transposed_output是转换后的结果。
- 参数
num_tensors – [in] 张量的数量。
input_list – [in] 2D 输入张量列表。
cast_output_list – [inout] 类型转换后的张量列表。 维度与 input_list 中的张量匹配。
transposed_output_list – [inout] 类型转换和转置后的张量列表。 维度是 input_list 中张量的转置。
stream – [in] 用于操作的 CUDA 流。
-
void nvte_dgeglu_cast_transpose(const NVTETensor input, const NVTETensor geglu_input, NVTETensor cast_output, NVTETensor transposed_output, cudaStream_t stream)¶
计算输入的 dgeglu，此外还对 dgeglu 输出进行类型转换和转置。
此函数产生 2 个结果
cast_output是转换的结果transposed_output是转换后的结果。
- 参数
input – [in] 形状为 [N, H] 的输入张量。
geglu_input – [in] 用作 GeGLU 操作正向传播输入的张量。形状 [N, H * 2]。
cast_output – [inout] 类型转换的结果。形状：[N, H * 2]。
transposed_output – [inout] 类型转换和转置的结果。形状：[H * 2, N]。
stream – [in] 用于操作的 CUDA 流。