MUTLASS

MUTLASS 是摩尔线程提供的高性能张量线性代数子程序库，专门针对 MUSA 架构进行了深度优化。

概述

MUTLASS 提供了 GEMM（矩阵乘法）及相关计算的高度优化实现。

什么是 MUTLASS

MUTLASS（MUSA Tensor Linear Algebra Subprogram Library，MUSA 张量线性代数子程序库）是一套 MUSA C++ 模板抽象集合，用于实现高性能矩阵乘法 (GEMM) 及相关计算。它包含分层分解和数据移动的策略，通过可重用的模块化软件组件实现。

关键特性

• 分层 GEMM: 完整的 GEMM 层次结构实现
• 混合精度: FP16、BF16、FP32、FP64、INT8
• Tensor Core: 支持 MUSA Tensor Core 指令
• Epilogue 融合: 激活函数、偏置添加等后处理融合
• MuTe（MUSA Templates for Tensors）布局: 强大的张量布局抽象

快速开始

安装源码

从 GitHub 获取 MUTLASS 源码：

git clone https://github.com/MooreThreads/mutlass.git

环境检查

使用 MUTLASS 前，请确认当前 MUSA 开发环境中包含 MUTLASS 头文件：

# 检查 MUTLASS 是否安装

ls $MUSA_HOME/include/mutlass/

# 编译选项

mcc -I$MUSA_HOME/include your_program.mu -L$MUSA_HOME/lib -lmutlass

简单 GEMM 示例

#include <mutlass/gemm/device/gemm_universal_adapter.h>
#include <mutlass/epilogue/thread/linear_combination.h>
#include <mutlass/util/host_tensor.h>

int main() {
    // 定义矩阵尺寸
    int M = 4096, N = 4096, K = 4096;
    
    // 创建张量 (使用 MuTe 布局)
    using LayoutA = mutlass::layout::RowMajor;
    using LayoutB = mutlass::layout::RowMajor;
    using LayoutC = mutlass::layout::RowMajor;
    
    // 分配设备内存
    float *d_A, *d_B, *d_C;
    musaMalloc(&d_A, M * K * sizeof(float));
    musaMalloc(&d_B, K * N * sizeof(float));
    musaMalloc(&d_C, M * N * sizeof(float));
    
    // 初始化输入矩阵
    init_matrix(d_A, M, K);
    init_matrix(d_B, K, N);
    
    // 使用简化的 GEMM 接口 (使用 MuTe 库的 gemm_device)
    // 注意: MUTLASS 没有简化接口，请使用完整的 GEMM API
    // 以下示例展示完整的 GEMM 调用方式
    
    // 1. 定义类型
    using ElementA = float;
    using ElementB = float;
    using ElementC = float;
    using ElementAccumulator = float;
    
    // 2. 构造 GEMM 参数
    // 注意: 需要使用实际的 GEMM kernel，这里仅作示例
    // 完整的 GEMM 使用需要参考 CollectiveBuilder 构建 kernel
    
    musaDeviceSynchronize();
    
    return 0;
}

GEMM 编程接口

核心 GEMM 概念

       A (M×K)         B (K×N)         C (M×N)
    ┌─────────┐      ┌─────────┐      ┌─────────┐
    │         │      │         │      │         │
    │  Tile   │  ×   │  Tile   │  =   │  Tile   │
    │         │      │         │      │         │
    └─────────┘      └─────────┘      └─────────┘
    
    CTA (Block) → Warp → Thread → MMA Instruction

完整 GEMM 示例 (Tensor Core)

#include <mutlass/mutlass.h>
#include <mutlass/gemm/device/gemm_universal_adapter.h>
#include <mutlass/gemm/collective/collective_builder.hpp>
#include <mutlass/epilogue/collective/collective_builder.hpp>
#include <mutlass/epilogue/thread/linear_combination.h>

int main() {
    // 1. 定义类型别名
    using ElementA = mutlass::half_t;      // FP16
    using ElementB = mutlass::half_t;
    using ElementC = float;                 // 累加用 FP32
    using ElementAccumulator = float;
    
    using LayoutA = mutlass::layout::RowMajor;
    using LayoutB = mutlass::layout::RowMajor;
    using LayoutC = mutlass::layout::RowMajor;
    
    // 2. 定义 GEMM kernel 配置
    using ArchTag = mutlass::arch::Mp31;    // MUSA MP31 架构
    using OperatorClass = mutlass::arch::OpClassTensorOp;
    using TileShape = mute::Shape<_128, _128, _64>;  // Threadblock tile 大小
    using ClusterShape = mute::Shape<_1, _1, _1>;     // Cluster 形状
    
    // 3. 使用 CollectiveBuilder 构建 mainloop
    using CollectiveMainloop = typename mutlass::gemm::collective::CollectiveBuilder<
        ArchTag, OperatorClass,
        ElementA, LayoutA, 8,
        ElementB, LayoutB, 8,
        ElementAccumulator,
        TileShape, ClusterShape,
        mutlass::gemm::collective::StageCount<4>,
        mutlass::gemm::KernelTmeWarpSpecialized
    >::CollectiveOp;
    
    // 4. 使用 CollectiveBuilder 构建 epilogue
    using ThreadEpilogueOp = mutlass::epilogue::fusion::LinearCombination<
        ElementC, float, ElementC, float>;
    
    using CollectiveEpilogue = typename mutlass::epilogue::collective::CollectiveBuilder<
        ArchTag, OperatorClass,
        TileShape, ClusterShape,
        mutlass::epilogue::collective::EpilogueTileAuto,
        ElementAccumulator, ElementAccumulator,
        ElementC, LayoutC, 8,
        ElementC, LayoutC, 8,
        mutlass::epilogue::NoSmem,
        ThreadEpilogueOp,
        CollectiveMainloop
    >::CollectiveOp;
    
    // 5. 定义 GEMM kernel
    using GemmKernel = mutlass::gemm::kernel::GemmUniversal<
        mute::Shape<int, int, int>,  // ProblemShape
        CollectiveMainloop,
        CollectiveEpilogue,
        mutlass::gemm::PersistentScheduler
    >;
    
    // 6. 创建 GEMM adapter
    using Gemm = mutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
    
    // 7. 构造 GEMM 参数
    typename Gemm::Arguments args(
        mutlass::gemm::GemmUniversalMode::kGemm,
        {4096, 4096, 4096},
        {(ElementA*)d_A, 4096},
        {(ElementB*)d_B, 4096},
        {{1.0f, 0.0f}, (ElementC*)d_C, 4096, (ElementC*)d_C, 4096}
    );
    
    // 8. 分配工作空间
    size_t workspace_size = Gemm::get_workspace_size(args);
    void* workspace;
    musaMalloc(&workspace, workspace_size);
    
    // 9. 创建 GEMM 操作对象
    Gemm gemm_op;
    auto status = gemm_op.initialize(args, workspace);
    MUTLASS_CHECK(status);
    
    // 10. 执行 GEMM
    status = gemm_op();
    MUTLASS_CHECK(status);
    
    musaFree(workspace);
    
    return 0;
}

批量 GEMM

// 批量矩阵乘法 - 使用 GemmUniversalAdapter 支持批量模式
// MUTLASS 使用 kGemmBatched 模式进行批量 GEMM

int batch_count = 32;

// 使用批量 GEMM 模式
typename Gemm::Arguments batch_args(
    mutlass::gemm::GemmUniversalMode::kGemmBatched,
    {M, N, K},
    {(ElementA*)d_A_array, stride_A},   // 批量指针数组
    {(ElementB*)d_B_array, stride_B},   // 批量指针数组
    {{alpha, beta}, (ElementC*)d_C_array, stride_C, (ElementC*)d_C_array, stride_C},
    batch_count
);

// 分配工作空间
size_t workspace_size = Gemm::get_workspace_size(batch_args);
void* workspace;
musaMalloc(&workspace, workspace_size);

// 创建 GEMM 操作
Gemm batch_gemm;
auto status = batch_gemm.initialize(batch_args, workspace);
MUTLASS_CHECK(status);

// 执行批量 GEMM
status = batch_gemm();
MUTLASS_CHECK(status);

布局 (Layout)

行主序和列主序

// 行主序 (RowMajor)
using LayoutRM = mutlass::layout::RowMajor;
LayoutRM layout_a(ld_a);  // ld_a = N

// 列主序 (ColumnMajor) - 相当于转置
using LayoutCM = mutlass::layout::ColumnMajor;
LayoutCM layout_b(ld_b);  // ld_b = M

复杂布局

// 平铺布局 (Tiled Layout)
using TiledLayout = mutlass::layout::Tile<
    mutlass::layout::RowMajor,
    128,   // tile M
    256    // tile N
>;

// 混合布局
using ComplexLayout = mutlass::layout::LayoutTransformer<
    LayoutA,
    mutlass::layout::Swizzle<4>  // swizzle 模式
>;

混合精度

FP16 输入，FP32 累加

using ElementA = mutlass::half_t;   // FP16 输入
using ElementB = mutlass::half_t;
using ElementC = mutlass::half_t;   // FP16 输出
using ElementAccum = float;          // FP32 累加器

INT8 量化

using ElementA = mutlass::int8_t;
using ElementB = mutlass::int8_t;
using ElementC = mutlass::int32_t;   // INT32 输出
using ElementAccum = mutlass::int32_t;

TF32 (Tensor Float 32)

using ElementA = mutlass::tfloat32_t;
using ElementB = mutlass::tfloat32_t;
using ElementC = mutlass::tfloat32_t;
using ElementAccum = float;

Epilogue 融合

偏置和激活

using Epilogue = mutlass::epilogue::thread::LinearCombination<
    ElementC,
    ElementAccumulator,
    ElementAccumulator,
    mutlass::epilogue::thread::ReLU  // ReLU 激活
>;

自定义 Epilogue

template <typename ElementOutput, typename ElementAccumulator>
class MyEpilogue {
public:
    using ElementCompute = ElementAccumulator;
    
    void operator()(
        ElementOutput* output,
        ElementAccumulator* accumulator,
        ElementAccumulator* source,
        ElementAccumulator* bias
    ) {
        // 自定义融合逻辑
        ElementAccumulator value = *source + *bias;
        value = fmaxf(value, 0);  // ReLU
        *output = (ElementOutput)value;
    }
};

性能优化

选择最佳内核

// 使用 Profiler 查找最佳配置
int main() {
    mutlass::profiler::OperationPerf KernalPerf;
    
    // 测试不同配置
    std::vector<mutlass::gemm::GemmCoord> problem_sizes = {
        {4096, 4096, 4096},
        {8192, 8192, 8192},
    };
    
    // 运行性能测试
    KernalPerf.profile(
        "gemm_kernel",
        problem_sizes,
        ElementA{}, ElementB{}, ElementC{}
    );
}

优化参数

// 调整块大小和 Warp 数量
constexpr int kStageCount = 3;     // 共享内存阶段数
constexpr int kWarpCountM = 4;       // M 方向的 Warp 数
constexpr int kWarpCountN = 2;       // N 方向的 Warp 数

// 选择数学指令
using InstructionShape = mutlass::gemm::GemmShape<16, 8, 16>;  // MMA 指令形状

MuTe 布局系统

MuTe 简介

MuTe 是 MUTLASS 3.x 引入的张量布局抽象库，提供层次化的多维布局表达能力。对应的 C++ 命名空间通常为 mute。

基本用法

#include <mutlass/tensor.h>
#include <mutlass/layout/layout.h>

// 创建 2D 布局
auto layout = mute::make_layout(
    mute::Shape<4096, 4096>{},    // 形状
    mute::Stride<4096, 1>{}       // 步长
);

// 索引转换
auto idx = layout(2, 3);    // 获取坐标 (2,3) 对应的线性索引

调试与诊断

错误检查

#define MUTLASS_CHECK(status)                                    \\
    do {                                                          \\
        auto err = status;                                        \\
        if (err != mutlass::Status::kSuccess) {                   \\
            std::cerr << \"mutlass Error: \" << (int)err << \"\\n\";  \\
            return err;                                           \\
        }                                                         \\
    } while (0)

// 使用
MUTLASS_CHECK(gemm.initialize(args));
MUTLASS_CHECK(gemm());

Profiler 使用

# 运行性能分析

$MUSA_HOME/bin/mutlass_profiler \
    --operation=gemm \
    --m=4096 --n=4096 --k=4096 \
    --A=f16:row --B=f16:row --C=f32:row

完整示例：神经网络前向传播

#include <mutlass/gemm/gemm.h>
#include <cutlass/epilogue/thread/linear_combination.h>

// Transformer 中的注意力机制示例
void attention_forward(
    float* Q, float* K, float* V, 
    float* output, 
    int batch, int seq_len, int heads, int dim
) {
    int M = batch * seq_len;
    int N = seq_len;
    int K_dim = dim / heads;  // 每个头的维度
    
    using Element = mutlass::half_t;
    
    // 使用 MUTLASS GEMM 需要通过 CollectiveBuilder 构建
    // 以下为概念性示例，实际使用需参考完整 API
    // 
    // 注意: 完整的 Transformer 注意力实现需要组合多个 GEMM 操作
    // Q × K^T 和 Attn × V 两个步骤
    
    // 缩放
    float scale = 1.0f / sqrtf((float)K_dim);
    
    // Softmax (需要额外处理)
    
    // Attention × V
    // ... 类似实现
}
    
    // Softmax (需要额外处理)
    
    // Attention × V
    // ... 类似实现
}

限制与注意事项

1. 架构兼容性: 需要支持 Tensor Core 的 MUSA 设备 (S4000/S5000)
2. 内存对齐: 建议使用 128 字节对齐以获得最佳性能
3. 工作空间: 大型 GEMM 操作需要额外的设备内存作为工作空间
4. 编译时间: 完整编译可能需要较长时间

版本信息

版本	日期
MULASS v0.3.0	Dec 19, 2025
MULASS v0.2.0	Feb 26, 2025
MULASS v0.1.1	Sep 30, 2024

相关文档

• MUSA OpenLibs