muBLASLt 开发者指南

muBLASLt 是 muBLAS 的扩展库，提供更高级的矩阵运算功能。相比基础 muBLAS，muBLASLt 支持更灵活的配置和更优的性能优化。

概述

什么是 muBLASLt

muBLASLt（Basic Linear Algebra Subroutines Lightweight based on MUSA，轻量级 GEMM 库）面向需要更细粒度矩阵乘法控制的场景，提供矩阵布局、算法选择、epilogue 融合和工作区配置等能力，帮助开发者在 MTGPU 上获得更灵活的 GEMM 性能优化空间。

关键特性

特性	说明
灵活矩阵布局	支持行主序/列主序/分块布局 (COL32, COL4_4R2_8C 等)
多种数据类型	FP16, BF16, TF32, FP32, FP64, FP8(E4M3/E5M2) 等
算法自动选择	通过启发式搜索自动选择最优算法
张量核心支持	利用 MT GPU 张量核心加速
工作区管理	支持自定义工作区内存

与 muBLAS 的对比

特性	muBLAS	muBLASLt
接口复杂度	简单	中等
性能优化	自动	可配置
数据类型	基础类型	扩展类型
算法选择	固定	启发式搜索

快速开始

以下示例展示如何使用 muBLASLt 执行 FP16 矩阵乘法，并融合 Bias + ReLU 激活：

#include <cstdio>
#include <cstdlib>
#include <vector>
#include <musa_runtime.h>
#include <mublasLt.h>
#include <musa_fp16.h>
#include <musa.h>

#define MUBLAS_CHECK(cmd)                                             \
  do {                                                                \
    mublasStatus_t status = (cmd);                                    \
    if (status != MUBLAS_STATUS_SUCCESS) {                            \
      fprintf(stderr, "muBLASLt error %d at %s:%d\n",                 \
              status, __FILE__, __LINE__);                            \
      exit(EXIT_FAILURE);                                             \
    }                                                                 \
  } while (0)

#define MUSA_CHECK(cmd)                                               \
  do {                                                                \
    musaError_t err = (cmd);                                          \
    if (err != musaSuccess) {                                         \
      fprintf(stderr, "MUSA error: %s at %s:%d\n",                    \
              musaGetErrorString(err), __FILE__, __LINE__);           \
      exit(EXIT_FAILURE);                                             \
    }                                                                 \
  } while (0)

int main() {
    const int M = 256, N = 256, K = 256;
    const __half alpha = __float2half(1.0f);
    const __half beta  = __float2half(0.0f);

    // ===== 1. 创建句柄 =====
    mublasLtHandle_t handle;
    MUBLAS_CHECK(mublasLtCreate(&handle));

    // ===== 2. 矩阵布局 =====
    mublasLtMatrixLayout_t Adesc, Bdesc, Cdesc, Ddesc;
    MUBLAS_CHECK(mublasLtMatrixLayoutCreate(&Adesc, MUSA_R_16F, M, K, K));
    MUBLAS_CHECK(mublasLtMatrixLayoutCreate(&Bdesc, MUSA_R_16F, K, N, N));
    MUBLAS_CHECK(mublasLtMatrixLayoutCreate(&Cdesc, MUSA_R_16F, M, N, N));
    MUBLAS_CHECK(mublasLtMatrixLayoutCreate(&Ddesc, MUSA_R_16F, M, N, N));

    // ===== 3. Matmul 描述符 =====
    mublasLtMatmulDesc_t computeDesc;
    MUBLAS_CHECK(mublasLtMatmulDescCreate(&computeDesc,
        MUBLAS_COMPUTE_16F, MUSA_R_16F));
    mublasOperation_t transa = MUBLAS_OP_N;
    mublasOperation_t transb = MUBLAS_OP_N;
    MUBLAS_CHECK(mublasLtMatmulDescSetAttribute(
        computeDesc, MUBLASLT_MATMUL_DESC_TRANSA,
        &transa, sizeof(transa)));
    MUBLAS_CHECK(mublasLtMatmulDescSetAttribute(
        computeDesc, MUBLASLT_MATMUL_DESC_TRANSB,
        &transb, sizeof(transb)));

    // ===== 4. Epilogue (RELU + BIAS) =====
    mublasLtEpilogue_t epilogue = MUBLASLT_EPILOGUE_RELU_BIAS;
    MUBLAS_CHECK(mublasLtMatmulDescSetAttribute(
        computeDesc, MUBLASLT_MATMUL_DESC_EPILOGUE,
        &epilogue, sizeof(epilogue)));

    // ===== 5. Preference =====
    mublasLtMatmulPreference_t preference;
    MUBLAS_CHECK(mublasLtMatmulPreferenceCreate(&preference));
    size_t workspaceSize = 32 * 1024 * 1024; // 32MB
    MUBLAS_CHECK(mublasLtMatmulPreferenceSetAttribute(
        preference,
        MUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
        &workspaceSize,
        sizeof(workspaceSize)));

    // ===== 6. heuristic =====
    mublasLtMatmulHeuristicResult_t heuristicResult;
    int returnedResults = 0;
    MUBLAS_CHECK(mublasLtMatmulAlgoGetHeuristic(
        handle, computeDesc,
        Adesc, Bdesc, Cdesc, Ddesc,
        preference, 1,
        &heuristicResult,
        &returnedResults));
    if (returnedResults == 0) {
        fprintf(stderr, "No algorithm found\n");
        return -1;
    }

    // ===== 7. 设备内存 =====
    __half *d_A, *d_B, *d_C, *d_D, *d_bias;
    void* workspace;
    MUSA_CHECK(musaMalloc(&d_A, M * K * sizeof(__half)));
    MUSA_CHECK(musaMalloc(&d_B, K * N * sizeof(__half)));
    MUSA_CHECK(musaMalloc(&d_C, M * N * sizeof(__half)));
    MUSA_CHECK(musaMalloc(&d_D, M * N * sizeof(__half)));
    MUSA_CHECK(musaMalloc(&d_bias, N * sizeof(__half)));
    MUSA_CHECK(musaMalloc(&workspace, workspaceSize));

    // ===== 8. Host 数据 =====
    std::vector<__half> h_A(M * K);
    std::vector<__half> h_B(K * N);
    std::vector<__half> h_C(M * N);
    std::vector<__half> h_bias(N);
    for (int i = 0; i < M * K; i++) h_A[i] = __float2half(1.0f);
    for (int i = 0; i < K * N; i++) h_B[i] = __float2half(1.0f);
    for (int i = 0; i < M * N; i++) h_C[i] = __float2half(0.0f);
    for (int i = 0; i < N; i++)     h_bias[i] = __float2half(0.5f);

    // ===== 9. 拷贝 =====
    MUSA_CHECK(musaMemcpy(d_A, h_A.data(), M*K*sizeof(__half), musaMemcpyHostToDevice));
    MUSA_CHECK(musaMemcpy(d_B, h_B.data(), K*N*sizeof(__half), musaMemcpyHostToDevice));
    MUSA_CHECK(musaMemcpy(d_C, h_C.data(), M*N*sizeof(__half), musaMemcpyHostToDevice));
    MUSA_CHECK(musaMemcpy(d_bias, h_bias.data(), N*sizeof(__half), musaMemcpyHostToDevice));

    // ===== 10. 设置 bias =====
    MUBLAS_CHECK(mublasLtMatmulDescSetAttribute(
        computeDesc,
        MUBLASLT_MATMUL_DESC_BIAS_POINTER,
        &d_bias,
        sizeof(d_bias)));

    // ===== 11. 执行 =====
    MUBLAS_CHECK(mublasLtMatmul(
        handle,
        computeDesc,
        &alpha,
        d_A, Adesc,
        d_B, Bdesc,
        &beta,
        d_C, Cdesc,
        d_D, Ddesc,
        &heuristicResult.algo,
        workspace,
        workspaceSize,
        0));
    MUSA_CHECK(musaStreamSynchronize(0));

    // ===== 12. 拿结果 =====
    std::vector<__half> h_D(M * N);
    MUSA_CHECK(musaMemcpy(h_D.data(), d_D, M*N*sizeof(__half), musaMemcpyDeviceToHost));
    float result = __half2float(h_D[0]);
    printf("D[0] = %f (expected ~ %f)\n", result, K + 0.5f);

    // ===== 13. 释放 =====
    MUSA_CHECK(musaFree(d_A));
    MUSA_CHECK(musaFree(d_B));
    MUSA_CHECK(musaFree(d_C));
    MUSA_CHECK(musaFree(d_D));
    MUSA_CHECK(musaFree(d_bias));
    MUSA_CHECK(musaFree(workspace));
    MUBLAS_CHECK(mublasLtMatmulPreferenceDestroy(preference));
    MUBLAS_CHECK(mublasLtMatmulDescDestroy(computeDesc));
    MUBLAS_CHECK(mublasLtMatrixLayoutDestroy(Adesc));
    MUBLAS_CHECK(mublasLtMatrixLayoutDestroy(Bdesc));
    MUBLAS_CHECK(mublasLtMatrixLayoutDestroy(Cdesc));
    MUBLAS_CHECK(mublasLtMatrixLayoutDestroy(Ddesc));
    MUBLAS_CHECK(mublasLtDestroy(handle));

    return 0;
}

编译命令:

mcc -x musa gemm_fp16_example.cpp \
    -L/usr/local/musa/lib \
    -I/usr/local/musa/include \
    -lmublasLt -lmublas -lmusart -lmusa \
    -o gemm_fp16_example

运行:

./gemm_fp16_example

API 参考

核心函数

函数	说明
`mublasLtCreate`	创建 muBLASLt 句柄
`mublasLtDestroy`	销毁 muBLASLt 句柄
`mublasLtMatmul`	执行矩阵乘法
`mublasLtMatrixTransform`	执行矩阵变换

矩阵布局

函数	说明
`mublasLtMatrixLayoutCreate`	创建矩阵布局描述符
`mublasLtMatrixLayoutInit`	初始化矩阵布局描述符
`mublasLtMatrixLayoutDestroy`	销毁矩阵布局描述符
`mublasLtMatrixLayoutSetAttribute`	设置矩阵布局属性

算法选择

函数	说明
`mublasLtMatmulAlgoGetHeuristic`	获取启发式算法
`mublasLtMatmulAlgoInit`	初始化算法描述符
`mublasLtMatmulAlgoCheck`	检查算法兼容性

性能优化建议

1. 使用合适的矩阵布局

// 列主序（默认）
mublasLtMatrixLayout_t desc;
const int rows = 8;
const int cols = 8;
mublasLtOrder_t order = MUBLASLT_ORDER_COL;
mublasLtMatrixLayoutCreate(&desc, MUSA_R_32F, rows, cols, rows);
mublasLtMatrixLayoutSetAttribute(desc, MUBLASLT_MATRIX_LAYOUT_ORDER,
    &order, sizeof(order));

// 行主序
order = MUBLASLT_ORDER_ROW;
mublasLtMatrixLayoutSetAttribute(desc, MUBLASLT_MATRIX_LAYOUT_ORDER,
    &order, sizeof(order));

// 分块布局（适用于张量核心）
order = MUBLASLT_ORDER_COL32;
mublasLtMatrixLayoutSetAttribute(desc, MUBLASLT_MATRIX_LAYOUT_ORDER,
    &order, sizeof(order));

mublasLtMatrixLayoutDestroy(desc);

2. 配置工作区大小

// 较大的工作区可以获得更好的性能
mublasLtMatmulPreference_t preference;
mublasLtMatmulPreferenceCreate(&preference);

size_t workspaceSize = 64 * 1024 * 1024;  // 64MB
mublasLtMatmulPreferenceSetAttribute(preference,
    MUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
    &workspaceSize, sizeof(workspaceSize));

mublasLtMatmulPreferenceDestroy(preference);

3. 使用 FP16/BF16 加速

// 使用半精度浮点
mublasLtMatmulDesc_t computeDesc;
mublasLtMatmulDescCreate(&computeDesc, MUBLAS_COMPUTE_16F, MUSA_R_16F);
mublasLtMatmulDescDestroy(computeDesc);

// 使用 Brain Float
mublasLtMatmulDescCreate(&computeDesc, MUBLAS_COMPUTE_32F, MUSA_R_16BF);
mublasLtMatmulDescDestroy(computeDesc);

数据类型支持

数据类型	说明	MUSA 数据类型
FP32	单精度浮点	`MUSA_R_32F`
FP64	双精度浮点	`MUSA_R_64F`
FP16	半精度浮点	`MUSA_R_16F`
BF16	Brain Float	`MUSA_R_16BF`
FP8 E4M3	8 位浮点 (4 指数 3 尾数)	`MUSA_R_8F_E4M3`
FP8 E5M2	8 位浮点 (5 指数 2 尾数)	`MUSA_R_8F_E5M2`

muBLASLt 开发者指南

概述

什么是 muBLASLt

关键特性

与 muBLAS 的对比

快速开始

API 参考

核心函数

矩阵布局

算法选择

性能优化建议

1. 使用合适的矩阵布局

2. 配置工作区大小

3. 使用 FP16/BF16 加速

数据类型支持

相关文档

这篇文章对您有帮助吗？

概述​

什么是 muBLASLt​

关键特性​

与 muBLAS 的对比​

快速开始​

API 参考​

核心函数​

矩阵布局​

算法选择​

性能优化建议​

1. 使用合适的矩阵布局​

2. 配置工作区大小​

3. 使用 FP16/BF16 加速​

数据类型支持​

相关文档​

这篇文章对您有帮助吗？

概述

什么是 muBLASLt

关键特性

与 muBLAS 的对比

快速开始

API 参考

核心函数

矩阵布局

算法选择

性能优化建议

1. 使用合适的矩阵布局

2. 配置工作区大小

3. 使用 FP16/BF16 加速

数据类型支持

相关文档