FlagBLAS 用户指南

FlagBLAS 用户指南#

使用 FlagBLAS#

FlagBLAS 直接与 PyTorch 集成。导入包并对 CUDA 张量调用算子:

import torch
import flag_blas

# 在 CUDA 上创建张量
a = torch.randn(1024, 1024, device='cuda')
b = torch.randn(1024, 1024, device='cuda')

# 矩阵乘法 (GEMM)
c = flag_blas.ops.sgemm(a, b)

算子列表#

完整的算子注册表维护在 FlagBLAS conf/operators.yaml