gemm.h
矩阵乘法函数。
函数
-
void nvte_cublas_gemm(const NVTETensor A, const NVTETensor B, NVTETensor D, const NVTETensor bias, NVTETensor pre_gelu_out, bool transa, bool transb, bool grad, NVTETensor workspace, bool accumulate, bool use_split_accumulator, int math_sm_count, cudaStream_t stream)
计算 2 个矩阵的矩阵乘法，可能与其他操作融合。
计算
D = AB如果bias和pre_gelu_out都是空张量D = AB + bias如果pre_gelu_out为空且bias不为空D = GELU(AB + bias)如果bias和pre_gelu_out都不是空张量
- 参数:
A – [in] A 矩阵。
B – [in] B 矩阵。
D – [inout] 输出矩阵。
bias – [in] 偏置张量。
pre_gelu_out – [inout] GELU 激活之前的输出矩阵。
transa – [in] A 矩阵是否转置。
transb – [in] B 矩阵是否转置。
grad – [in] 此操作是否为梯度计算的一部分。
workspace – [out] 工作区张量。
accumulate – [in] 是否将结果累积到 D 矩阵中。
use_split_accumulator – [in] 是否在 FP8 GEMM 中使用拆分累加器。
math_sm_count – [in] 要使用的 GPU SM 数量 (默认=0: 使用 cuBLAS 启发式方法)
stream – [in] 用于操作的 CUDA 流。
-
void nvte_cublas_atomic_gemm(const NVTETensor A, const NVTETensor B, NVTETensor D, const NVTETensor bias, NVTETensor pre_gelu_out, bool transa, bool transb, bool grad, NVTETensor workspace, bool accumulate, bool use_split_accumulator, int math_sm_count, int m_split, int n_split, bool gemm_producer, const NVTETensor counter, cudaStream_t stream)
计算具有分块和原子计数器的 2 个矩阵的矩阵乘法。
计算
D = AB如果bias和pre_gelu_out都是空张量D = AB + bias如果pre_gelu_out为空且bias不为空D = GELU(AB + bias)如果bias和pre_gelu_out都不是空张量
警告
Cublas 原子 gemm 使用 beta API，并且未针对所有用例进行测试。
- 参数:
A – [in] A 矩阵。
B – [in] B 矩阵。
D – [inout] 输出矩阵。
bias – [in] 偏置张量。
pre_gelu_out – [inout] GELU 激活之前的输出矩阵。
transa – [in] A 矩阵是否转置。
transb – [in] B 矩阵是否转置。
grad – [in] 此操作是否为梯度计算的一部分。
workspace – [out] 工作区张量。
accumulate – [in] 是否将结果累积到 D 矩阵中。
use_split_accumulator – [in] 是否在 FP8 GEMM 中使用拆分累加器。
math_sm_count – [in] 要使用的 GPU SM 数量 (默认=0: 使用 cuBLAS 启发式方法)
m_split – [in] 沿 m 维度进行原子 GEMM 的块/分割数。
n_split – [in] 沿 n 维度进行原子 GEMM 的块/分割数。
gemm_producer – [in] 原子 GEMM 是生产者还是消费者。
counter – [inout] counter[chunk_i]=0 表示 chunk_i 已生成。
stream – [in] 用于操作的 CUDA 流。
-
void nvte_multi_stream_cublas_gemm(const NVTETensor *A, const NVTETensor *B, NVTETensor *D, const NVTETensor *bias, NVTETensor *pre_gelu_out, const int num_gemms, bool transa, bool transb, bool grad, NVTETensor *workspace, bool accumulate, bool use_split_accumulator, int math_sm_count, cudaStream_t stream)
计算多对矩阵乘法，可能与其他操作融合，在多个流上进行。
计算
D = AB如果bias和pre_gelu_out都是空张量D = AB + bias如果pre_gelu_out为空且bias不为空D = GELU(AB + bias)如果bias和pre_gelu_out都不是空张量
- 参数:
A – [in] A 矩阵列表。
B – [in] B 矩阵列表。
D – [inout] 输出矩阵列表。
bias – [in] 偏置张量列表。
pre_gelu_out – [inout] GELU 激活之前的输出矩阵列表。
num_gemms – [in] 要计算的 GEMM 数量。
transa – [in] A 矩阵是否转置。
transb – [in] B 矩阵是否转置。
grad – [in] 此操作是否为梯度计算的一部分。
workspace – [out] 工作区张量列表。
accumulate – [in] 是否将结果累积到 D 矩阵中。
use_split_accumulator – [in] 是否在 FP8 GEMM 中使用拆分累加器。
math_sm_count – [in] 要使用的 GPU SM 数量 (默认=0: 使用 cuBLAS 启发式方法)
stream – [in] 要等待的 CUDA 流。