跳到主要内容

MUTLASS 代码使用指南

目录​

概述​

MUTLASS (MUSA Templates for Linear Algebra Subroutines) 是一个用于在 MUSA 架构上实现高性能矩阵乘法运算的 C++ 模板库。本文档详细介绍如何在代码中使用 MUTLASS。

  • 纯头文件库：MUTLASS 是 header-only 库，只需包含头文件即可使用
  • 模板元编程：使用 C++ 模板实现零开销抽象
  • 分层设计：从底层原子操作到高级 API 的多层抽象
  • 架构优化：针对 MUSA 架构（Quyuan/MP22, MP31）优化

环境配置​

1. 包含路径设置​

需要将 MUTLASS 的 include/ 目录添加到包含路径：

# CMakeLists.txt
target_include_directories(target PRIVATE
    ${MUTLASS_DIR}/include
)

或者使用编译器选项：

mcc -I/path/to/mutlass/include your_code.mu

2. 编译要求​

  • C++ 标准：C++17 或更高
  • 编译器：MCC 4.3.4+
  • 文件扩展名：.mu 文件自动启用 MUSA 扩展，.cpp/.cc 文件需要 -x musa 选项
# .mu 文件
mcc code.mu -std=c++17

# .cpp 文件
mcc -x musa code.cpp -std=c++17

3. 链接库​

如果使用 MUTLASS Library（预编译的 kernel 实例），需要链接：

target_link_libraries(target PRIVATE
    mutlass_lib
    musart
    musa_driver
)

完整项目构建与运行​

本节详细介绍如何编译和运行 MUTLASS 项目的所有组件，包括 Examples、Tools、Tests 和 Experimental 功能。建议新用户先阅读本节，快速上手项目。

1. 完整项目构建​

1.1 清理和准备​

# 进入项目目录
cd /home/test/mutlass

# 清理旧的构建（如果存在）
rm -rf build

# 创建新的构建目录
mkdir build && cd build

1.2 CMake 配置​

完整配置（推荐用于开发）：

S4000将-DMUTLASS_MCC_ARCHS=31设为-DMUTLASS_MCC_ARCHS=22

cmake .. \
    -DMUTLASS_MCC_ARCHS=31 \
    -DMUTLASS_ENABLE_EXAMPLES=ON \
    -DMUTLASS_ENABLE_TESTS=ON \
    -DMUTLASS_ENABLE_PROFILER=ON \
    -DMUTLASS_ENABLE_LIBRARY=ON \
    -DMUTLASS_ENABLE_EXPERIMENTAL=ON

警告修复：（可忽略）将cmake/googletest.cmake中line48-54替换为

if(NOT googletest_POPULATED)
  if (MSVC)
    set(gtest_force_shared_crt ON CACHE BOOL "" FORCE)
  endif()
  FetchContent_MakeAvailable(googletest)
endif()

最小配置（快速验证）：

cmake .. \
    -DMUTLASS_MCC_ARCHS=31 \
    -DMUTLASS_ENABLE_EXAMPLES=ON \
    -DMUTLASS_ENABLE_PROFILER=ON

多架构配置：

cmake .. \
    -DMUTLASS_MCC_ARCHS="22;31" \
    -DMUTLASS_ENABLE_EXAMPLES=ON \
    -DMUTLASS_ENABLE_PROFILER=ON

1.3 编译组件​

1.3.1整体编译​
# 编译所有目标
make -j$(nproc)

# 使用较少并行数避免依赖问题
make -j4
1.3.2分步骤编译（推荐）​
make mutlass_library_objs -j1  
make mutlass_library -j1      # Library（会自动编译 mutlass_library_objs）
make mutlass_library_static -j$(nproc) # 编译静态库
make mutlass_examples -j$(nproc)      # Examples
make mutlass_profiler -j$(nproc)      # Profiler，耗时较长
make mutlass_test_unit -j$(nproc)     # Tests
make mutlass_experimental -j$(nproc)  # Experimental

说明：

  • make mutlass_library 会自动构建大量子目标，包括：
    • 各种架构的 GEMM kernel（如 mutlass_library_gemm_mp22_sgemm_simt、mutlass_library_gemm_mp31_void_s128x64x64gemm_f16_tensorop 等）
    • 各种数据类型组合（sgemm, f16, e4m3, e5m2 等）
    • 各种 kernel 类型（simt, tensorop 等）
    • 每个 kernel 都有对应的 *_objs（对象文件）、共享库（.so）和静态库（.a）

此时核心组件编译完成

2. 设置运行环境​

运行 Profiler 或依赖 Library 的可执行程序前，需设置动态库路径：

# 在 build 目录下设置库路径（必需）
export LD_LIBRARY_PATH=$PWD/tools/library:$LD_LIBRARY_PATH

# 或使用绝对路径
export LD_LIBRARY_PATH=/home/test/mutlass/build/tools/library:$LD_LIBRARY_PATH

# 验证
echo $LD_LIBRARY_PATH

说明：Examples（如 00_basic_gemm）不依赖 libmutlass.so，可直接运行；Profiler 必须设置 LD_LIBRARY_PATH 才能找到 libmutlass.so。

3. 运行 Examples（示例代码）​

3.1 运行单个示例​


# 基础 GEMM 示例
./examples/00_basic_gemm/00_basic_gemm --m=2048 --n=2048 --k=1024

# Quyuan 架构集体构建器示例（需要 MP22 架构）
./examples/01_quyuan_gemm_with_collective_builder/01_collective_builder \
    --m=2048 --n=2048 --k=1024 --l=2

# MP31 FP8 GEMM 示例（需要 MP31 架构）
./examples/02_mp31_fp8_gemm_with_collective_builder/02_mp31_fp8_gemm

# MP31 FP8 缩放 GEMM 示例
./examples/03_mp31_fp8_scaling_gemm/03_fp8_scaling_gemm

3.2 运行所有示例​

# 运行所有示例测试（会自动编译并运行）
make test_examples

4. 运行 Profiler（性能分析工具）​

4.1 编译 Profiler​

cd /home/test/mutlass/build

# 编译 Profiler（需要先编译 Library）
make mutlass_profiler -j$(nproc)

4.2 运行 Profiler​

基础 GEMM 性能测试：


./tools/profiler/mutlass_profiler \
    --operation=Gemm \
    --op_class=simt \
    --m=4096 \
    --n=4096 \
    --k=1024 \
    --cta_m=128 \
    --cta_n=128

TensorOp GEMM 性能测试：

./tools/profiler/mutlass_profiler \
    --operation=Gemm \
    --op_class=tensorop \
    --m=4096 \
    --n=4096 \
    --k=1024

FP16 GEMM 性能测试：

./tools/profiler/mutlass_profiler \
    --operation=Gemm \
    --op_class=tensorop \
    --m=4096 \
    --n=4096 \
    --k=1024 \
    --A=f16:row \
    --B=f16:column \
    --C=f16:column \
    --D=f16:column

查看帮助信息：

# 通用帮助
./tools/profiler/mutlass_profiler --help

# GEMM 操作帮助
./tools/profiler/mutlass_profiler --operation=Gemm --help

Profiler 输出示例：

=============================
  Problem ID: 1

        Provider: MUTLASS
   OperationKind: gemm
       Operation: mutlass_mp22_simt_sgemm_f32_f32_f32_f32_f32_128x128x4_tnn_align2

          Status: Success
    Verification: ON
     Disposition: Passed

reference_device: Passed
          muBLAS: Not run
           muDNN: Not run

       Arguments: --gemm_kind=universal --m=4096 --n=4096 --k=1024 --A=f32:row --B=f32:column --C=f32:column --D=f32:column  \
                  --alpha=1 --beta=0 --batch_count=1 --op_class=simt --accum=f32 --cta_m=128 --cta_n=128 --cta_k=4 --cluster_m=1  \
                  --cluster_n=1 --cluster_k=1 --stages=2 --inst_m=1 --inst_n=1 --inst_k=1 --min_cc=22 --max_cc=1024

           Bytes: 100663296  bytes
           FLOPs: 34393292800  flops
           FLOPs/Byte: 341

         Runtime: 2.92589  ms
          Memory: 32.0416 GiB/s

            Math: 11754.8 GFLOP/s

5. 运行单元测试（Unit Tests）​

5.1 编译单元测试​

cd /home/test/mutlass/build

# 编译所有单元测试
make test_unit -j$(nproc)

# 或编译特定测试
make test_unit_mute_core
make test_unit_gemm

5.2 运行单元测试​

# 运行所有单元测试（会自动编译并运行）
make test_unit

# 运行特定测试组
make test_unit_mute_core
make test_unit_gemm_device

预期输出示例：

[==========] Running 33 tests from 1 test suite.
[----------] Global test environment set-up.
[ RUN      ] MuTe_core.Tuple
[       OK ] MuTe_core.Tuple (0 ms)
...
[----------] 33 tests from MuTe_core (7 ms total)
[==========] 33 tests from 1 test suite ran. (7 ms total)
[  PASSED  ] 33 tests.

6. 运行 Experimental（实验性功能）​

6.1 编译 Experimental​

cd /home/test/mutlass/build

# 编译所有实验性功能
make mutlass_experimental -j$(nproc)

包含的实验性功能：

  • mp31_fmha_fwd: MP31 架构的 Fused Multi-Head Attention (FMHA)
  • mp31_paged_fmha_fwd: MP31 架构的 Paged FMHA
  • mp31_mla_decode: MP31 架构的 Multi-Layer Attention (MLA)

6.2 运行 FMHA 测试​


# 运行所有 FMHA 测试
make test_examples_mp31_fmha_fwd

# 运行特定测试用例
make test_examples_mp31_fmha_fwd_test_causal_00_128_128
make test_examples_mp31_fmha_fwd_test_non_causal_00_128_128
make test_examples_mp31_fmha_fwd_test_varlen_causal_00_128_128

6.3 运行 Paged FMHA 测试​

# 运行所有 Paged FMHA 测试
make test_examples_mp31_paged_fmha_fwd

# 运行特定测试用例
make test_examples_mp31_paged_fmha_fwd_test_causal_00_128_128
make test_examples_mp31_paged_fmha_fwd_test_non_causal_00_128_128

6.4 运行 MLA 测试​

# 运行 MLA 测试
make test_examples_mp31_mla_decode

注意：实验性功能可能不稳定，API 可能会变化。

7. 运行所有测试（一键测试）​

cd /home/test/mutlass/build

# 运行所有测试（包括 examples, tests, experimental）
make test_all

这会依次运行：

  1. test_examples - 所有示例测试
  2. test_unit - 所有单元测试
  3. test_experimental - 所有实验性功能测试

8. 验证构建完整性​

8.1 检查编译目标​

cd /home/test/mutlass/build

# 查看所有可用目标
make help | grep -E '^\.\.\. (test_|mutlass_)' | head -50

# 统计目标数量
make help | grep -E '^\.\.\. (test_|mutlass_)' | wc -l

8.2 检查生成的文件​

# 检查可执行文件
find examples -type f -executable
find tools -type f -executable
find experimental -type f -executable

# 检查库文件
ls -lh tools/library/libmutlass*.so | head -10

# 检查库文件数量
ls tools/library/libmutlass*.so | wc -l

8.3 验证库依赖​

# 检查 profiler 的依赖
ldd tools/profiler/mutlass_profiler | grep mutlass

# 检查示例的依赖
ldd examples/00_basic_gemm/00_basic_gemm | grep musa

9. 完整验证脚本​

创建一个完整的验证脚本 verify_all.sh：

#!/bin/bash
set -e

echo "=== MUTLASS 完整验证脚本 ==="

# 1. 清理旧构建
echo "[1/7] 清理旧构建..."
cd /home/test/mutlass
rm -rf build

# 2. 配置 CMake
echo "[2/7] 配置 CMake..."
mkdir build && cd build
cmake .. \
    -DMUTLASS_MCC_ARCHS=31 \
    -DMUTLASS_ENABLE_EXAMPLES=ON \
    -DMUTLASS_ENABLE_TESTS=ON \
    -DMUTLASS_ENABLE_PROFILER=ON \
    -DMUTLASS_ENABLE_EXPERIMENTAL=ON

# 3. 编译所有组件
echo "[3/7] 编译所有组件..."
make -j$(nproc)

# 4. 设置环境
echo "[4/7] 设置环境变量..."
export LD_LIBRARY_PATH=$PWD/tools/library:$LD_LIBRARY_PATH

# 5. 运行示例测试
echo "[5/7] 运行示例测试..."
make test_examples

# 6. 运行 Profiler
echo "[6/7] 运行 Profiler..."
./tools/profiler/mutlass_profiler \
    --operation=Gemm \
    --op_class=simt \
    --m=1024 --n=1024 --k=512 \
    --cta_m=128 --cta_n=128 \
    > /dev/null 2>&1 && echo "Profiler 运行成功" || echo "Profiler 运行失败"

# 7. 运行单元测试
echo "[7/7] 运行单元测试..."
make test_unit

echo "=== 验证完成 ==="

运行验证脚本：

chmod +x verify_all.sh
./verify_all.sh

10. 常见运行问题​

10.1 找不到库文件​

错误：./tools/profiler/mutlass_profiler: error while loading shared libraries: libmutlass.so: cannot open shared object file: No such file or directorys

原因：通常是因为未编译 Library 组件，或者 LD_LIBRARY_PATH 未正确设置。

解决方法：

  1. 确保已编译 Library：
cd /home/test/mutlass/build
make mutlass_library -j1
  1. 设置库路径：
export LD_LIBRARY_PATH=/home/test/mutlass/build/tools/library:$LD_LIBRARY_PATH
  1. 验证库文件是否存在：
ls -lh tools/library/libmutlass*.so

如果库文件不存在，说明 Library 未编译成功，需要重新编译。

10.2 架构不匹配​

错误：某些示例需要特定架构（如 MP22 或 MP31）

解决：

  • 检查当前 GPU 架构：musaInfo
  • 使用匹配的架构编译：-DMUTLASS_MCC_ARCHS=22 或 31

10.3 测试超时​

错误：某些测试运行时间过长

解决：

  • 减少问题规模
  • 使用更少的并行数编译：make -j4
  • 单独运行测试而不是 make test_all

11. 性能基准测试​

11.1 GEMM 性能测试矩阵​

# 小规模测试
./tools/profiler/mutlass_profiler \
    --operation=Gemm --op_class=tensorop \
    --m=1024 --n=1024 --k=1024

# 中等规模测试
./tools/profiler/mutlass_profiler \
    --operation=Gemm --op_class=tensorop \
    --m=4096 --n=4096 --k=4096

# 大规模测试
./tools/profiler/mutlass_profiler \
    --operation=Gemm --op_class=tensorop \
    --m=8192 --n=8192 --k=8192

11.2 批量测试脚本​

#!/bin/bash
# benchmark.sh - 批量性能测试

export LD_LIBRARY_PATH=/home/test/mutlass/build/tools/library:$LD_LIBRARY_PATH

sizes=(1024 2048 4096 8192)
for size in "${sizes[@]}"; do
    echo "Testing size: ${size}x${size}x${size}"
    ./tools/profiler/mutlass_profiler \
        --operation=Gemm --op_class=tensorop \
        --m=$size --n=$size --k=$size \
        | grep -E "Runtime|Math|Memory"
done

12. 项目组件总结​

组件编译目标测试目标说明
Examplesmutlass_examplestest_examples4 个示例程序
Profilermutlass_profilertest_profiler性能分析工具
Librarymutlass_lib-预编译 kernel 库
Testsmutlass_test_unittest_unit单元测试
Experimentalmutlass_experimentaltest_experimental实验性功能

13. 快速参考命令​

# === 构建 ===
cd /home/test/mutlass && rm -rf build && mkdir build && cd build
cmake .. -DMUTLASS_MCC_ARCHS=31 -DMUTLASS_ENABLE_EXAMPLES=ON

# 容易多次error，必要时降低并行数
make -j$(nproc)

# === 运行测试 ===
make test_examples          # 示例测试
make test_unit              # 单元测试
make test_experimental      # 实验性功能测试
make test_all               # 所有测试

# === 运行工具 ===
./tools/profiler/mutlass_profiler --operation=Gemm --help

# === 运行示例 ===
./examples/00_basic_gemm/00_basic_gemm --m=2048 --n=2048 --k=1024

基础 API 使用示例​

1. 基本头文件包含示例​

以下头文件路径与 examples/、include/mutlass/、tools/util/include/mutlass/ 一致（构建时需将对应 include 目录加入编译路径）：

// 核心与数值类型
#include "mute/tensor.hpp"
#include "mutlass/mutlass.h"
#include "mutlass/numeric_types.h"
#include "mutlass/half.h"
#include "mutlass/bfloat16.h"
#include "mutlass/float8.h"

// GEMM：Builder 或低层 API
#include "mutlass/gemm/device/gemm_universal_adapter.h"
#include "mutlass/gemm/collective/collective_builder.hpp"
#include "mutlass/epilogue/collective/collective_builder.hpp"
#include "mutlass/gemm/dispatch_policy.hpp"

// 工具：设备/主机内存、stride、参考实现
#include "mutlass/util/device_memory.h"
#include "mutlass/util/host_tensor.h"
#include "mutlass/util/packed_stride.hpp"
#include "mutlass/util/reference/device/tensor_fill.h"
#include "mutlass/util/reference/device/tensor_compare.h"
#include "mutlass/util/reference/device/gett.hpp"

2. 数据类型使用​

#include "mutlass/numeric_types.h"
#include "mutlass/half.h"
#include "mutlass/bfloat16.h"
#include "mutlass/complex.h"

using namespace mutlass;

// 基本类型（32 位浮点用 float，16 位 half 用 half_t，bfloat16 用 bfloat16_t）
float f32_value = 1.0f;
half_t f16_value = 2.0_hf;       // 字面量后缀 _hf
bfloat16_t bf16_value = 3.0_bf16;  // 字面量后缀 _bf16

// 复数类型
complex<float> c32_value(1.0f, 2.0f);

3. 内存分配​

#include "mutlass/util/device_memory.h"

using namespace mutlass;

// 设备内存分配
DeviceAllocation<float> device_A(M * K);  // M x K 矩阵
DeviceAllocation<float> device_B(K * N);  // K x N 矩阵
DeviceAllocation<float> device_C(M * N);  // M x N 矩阵

// 主机内存分配
HostTensor<float> host_A({M, K});
HostTensor<float> host_B({K, N});
HostTensor<float> host_C({M, N});

// 数据拷贝
device_A.copy_from(host_A.data());
device_B.copy_from(host_B.data());

4. 使用 Collective Builder（推荐）​

Collective Builder 提供简化的 API，自动选择最优的模板参数。完整参数与调用方式请参考 examples/01_quyuan_gemm_with_collective_builder/collective_builder.mu。

#include "mutlass/gemm/collective/collective_builder.hpp"
#include "mutlass/epilogue/collective/collective_builder.hpp"
#include "mutlass/gemm/device/gemm_universal_adapter.h"
#include "mute/tensor.hpp"

using namespace mutlass;
using namespace mute;

// 类型与布局（需与架构匹配，如 MP22 用 arch::Mp22）
using ElementA = float;
using ElementB = float;
using ElementC = float;
using ElementD = float;
using ElementAccumulator = float;
using LayoutA = layout::ColumnMajor;
using LayoutB = layout::RowMajor;
using LayoutC = layout::ColumnMajor;
using LayoutD = layout::ColumnMajor;

using ArchTag = arch::Mp31;  // 或 arch::Mp22
using OpClass = arch::OpClassTensorOp;
static constexpr int AlignmentA = 16 / sizeof(ElementA);
static constexpr int AlignmentB = 16 / sizeof(ElementB);
static constexpr int AlignmentC = 16 / sizeof(ElementC);
static constexpr int AlignmentD = 16 / sizeof(ElementD);

// 使用 Builder 创建 Mainloop（需传入 ArchTag, OpClass, 对齐等）
using CollectiveMainloop = typename gemm::collective::CollectiveBuilder<
    ArchTag, OpClass,
    ElementA, LayoutA, AlignmentA,
    ElementB, LayoutB, AlignmentB,
    ElementAccumulator,
    Shape<_128, _128, _32>,
    Shape<_1, _1, _1>,
    gemm::collective::StageCountAuto,
    gemm::collective::KernelScheduleAuto
>::CollectiveOp;

// 使用 Builder 创建 Epilogue
using CollectiveEpilogue = typename epilogue::collective::CollectiveBuilder<
    ArchTag, OpClass,
    Shape<_128, _128, _32>,
    Shape<_1, _1, _1>,
    epilogue::collective::EpilogueTileAuto,
    ElementAccumulator, float,
    ElementC, LayoutC, AlignmentC,
    ElementD, LayoutD, AlignmentD,
    epilogue::collective::EpilogueScheduleAuto
>::CollectiveOp;

// 组合为 Kernel 并用 GemmUniversalAdapter 封装
using GemmKernel = gemm::kernel::GemmUniversal<
    Shape<int, int, int, int>,
    CollectiveMainloop,
    CollectiveEpilogue
>;
using Gemm = gemm::device::GemmUniversalAdapter<GemmKernel>;

// 实际调用需构造 Arguments、allocate workspace、initialize、run，见示例。

5. 使用更低层次 API​

如需更精细控制，可直接使用 CollectiveMma、CollectiveEpilogue 等。模板参数较多，建议以 examples/00_basic_gemm/basic_gemm.mu 或 collective builder 示例为参考。

#include "mutlass/gemm/collective/collective_mma.hpp"
#include "mutlass/epilogue/collective/collective_epilogue.hpp"
#include "mutlass/gemm/device/gemm_universal_adapter.h"

// CollectiveMma / CollectiveEpilogue 需指定完整的 TiledMma、SmemSwizzle 等类型，
// 具体以 include 下的 collective_mma.hpp、collective_epilogue.hpp 及示例为准。

// 组合成 Kernel 后，使用 gemm::kernel::GemmUniversal + GemmUniversalAdapter 封装并调用。

6. GEMM 调用流程（与 examples 一致）​

无论使用 Collective Builder 还是低层 API，设备侧 GEMM 的调用顺序一致（参见 00_basic_gemm/basic_gemm.mu、01_quyuan_gemm_with_collective_builder/collective_builder.mu）：

using Gemm = mutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;

// 1. 问题规模与 stride（packed 布局可用 make_mute_packed_stride）
ProblemShapeType problem_size{M, N, K, L};
StrideA stride_A = mutlass::make_mute_packed_stride(StrideA{}, mute::make_shape(M, K, L));
StrideB stride_B = ...;
StrideC stride_C = ...;
StrideD stride_D = ...;

// 2. 构造 Arguments（mode、problem_size、mainloop 指针与 stride、epilogue alpha/beta 与 C/D、KernelHardwareInfo）
typename Gemm::Arguments arguments{
    mutlass::gemm::GemmUniversalMode::kGemm,
    problem_size,
    {ptr_A, stride_A, ptr_B, stride_B},
    {{alpha, beta}, ptr_C, stride_C, ptr_D, stride_D},
    mutlass::KernelHardwareInfo{}
};

// 3. 分配 workspace 并执行
Gemm gemm;
size_t workspace_size = Gemm::get_workspace_size(arguments);
mutlass::device_memory::allocation<uint8_t> workspace(workspace_size);

MUTLASS_CHECK(gemm.can_implement(arguments));
MUTLASS_CHECK(gemm.initialize(arguments, workspace.get()));
MUTLASS_CHECK(gemm.run());

代码示例详解​

示例 1：基础 GEMM（00_basic_gemm）​

基于低层 API 的 SGEMM，面向 Quyuan（MP22）架构，使用 SIMT FMA。源码：examples/00_basic_gemm/basic_gemm.mu。

要点：

  • 布局与类型：LayoutA/B/C = layout::RowMajor，TypeA/B/C = float。
  • TiledMma：TiledMMA<MMA_Atom<UniversalFMA<TypeC,TypeA,TypeB>>, Layout<Shape<_16,_16,_1>>>。
  • TileShape：Shape<_128, _128, _4>；DispatchPolicy：MainloopMp22TwoStage。
  • Gmem/Smem：GmemTiledCopyA/B、SmemLayoutAtomA/B、SmemCopyAtomA/B 由 Copy atom 与线程布局组成。
  • CollectiveMma：gemm::collective::CollectiveMma<DispatchPolicy, TileShape, TypeA, TagToStrideA_t<LayoutA>, TypeB, TagToStrideB_t<LayoutB>, TiledMma, GmemTiledCopyA, SmemLayoutAtomA, SmemCopyAtomA, ..., GmemTiledCopyB, ...>。
  • Epilogue：epilogue::collective::DefaultEpilogue<..., LinearCombination<TypeC, 1, TypeC, TypeC>, gemm::EpilogueDefault>。
  • Kernel：GemmKernel = gemm::kernel::GemmUniversal<Shape<int,int,int,int>, CollectiveMainloop, CollectiveEpilogue>，Gemm = gemm::device::GemmUniversalAdapter<GemmKernel>。
  • 调用：构造 Arguments（GemmUniversalMode::kGemm、problem_size、stride、alpha/beta、C/D 指针），get_workspace_size → device_memory::allocation<uint8_t> → can_implement → initialize → run。

源码对照（节选自 basic_gemm.mu）：

// 布局与类型
using LayoutA = mutlass::layout::RowMajor;
using LayoutB = mutlass::layout::RowMajor;
using LayoutC = mutlass::layout::RowMajor;
using TypeA = float;
using TypeB = float;
using TypeC = float;

// TiledMma：FMA 原子 + 16x16x1 线程布局
using TiledMma = TiledMMA<
    MMA_Atom<UniversalFMA<TypeC, TypeA, TypeB>>,
    Layout<Shape<_16, _16, _1>>
>;
using TileShape = Shape<_128, _128, _4>;
using DispatchPolicy = mutlass::gemm::MainloopMp22TwoStage;

// CollectiveMma / CollectiveEpilogue / Kernel
using CollectiveMainloop = mutlass::gemm::collective::CollectiveMma<
  DispatchPolicy, TileShape,
  TypeA, mutlass::detail::TagToStrideA_t<LayoutA>,
  TypeB, mutlass::detail::TagToStrideB_t<LayoutB>,
  TiledMma,
  GmemTiledCopyA, SmemLayoutAtomA, SmemCopyAtomA, mute::identity,
  GmemTiledCopyB, SmemLayoutAtomB, SmemCopyAtomB, mute::identity
>;
using CollectiveEpilogue = mutlass::epilogue::collective::DefaultEpilogue<
  mutlass::detail::TagToStrideC_t<LayoutC>,
  mutlass::detail::TagToStrideC_t<LayoutC>,
  mutlass::epilogue::thread::LinearCombination<TypeC, 1, TypeC, TypeC>,
  mutlass::gemm::EpilogueDefault
>;
using GemmKernel = mutlass::gemm::kernel::GemmUniversal<
  Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue
>;
using Gemm = mutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;

// 调用流程：stride → Arguments → workspace → can_implement → initialize → run
StrideA stride_a = mutlass::make_mute_packed_stride(StrideA{}, mute::make_shape(M, K, Batch));
auto arguments = typename Gemm::Arguments{
  mutlass::gemm::GemmUniversalMode::kGemm, problem_size,
  {A, stride_a, B, stride_b},
  {{alpha, beta}, C, stride_c, C, stride_c},
  mutlass::KernelHardwareInfo{}
};
mutlass::device_memory::allocation<uint8_t> workspace(Gemm::get_workspace_size(arguments));
MUTLASS_CHECK(gemm.can_implement(arguments));
MUTLASS_CHECK(gemm.initialize(arguments, workspace.get()));
MUTLASS_CHECK(gemm.run());

示例 2：Collective Builder（01_quyuan_gemm_with_collective_builder）​

使用 Collective Builder 在 Quyuan（MP22）上组 FP16 GEMM，由 Builder 自动选择 MMA、Smem、调度等。源码：examples/01_quyuan_gemm_with_collective_builder/collective_builder.mu。

要点：

  • 类型：ElementA/B/C/D = half_t，ElementAccumulator = float；布局：LayoutA = ColumnMajor，LayoutB = RowMajor，LayoutC/D = ColumnMajor。
  • 架构：ArchTag = arch::Mp22，OpClass = arch::OpClassTensorOp；对齐：AlignmentA/B/C/D = 16 / sizeof(Element)。
  • Mainloop：gemm::collective::CollectiveBuilder<ArchTag, OpClass, ElementA, LayoutA, AlignmentA, ElementB, LayoutB, AlignmentB, ElementAccumulator, Shape<_256,_128,_32>, Shape<_1,_1,_1>, StageCountAuto, KernelScheduleAuto>::CollectiveOp。
  • Epilogue：epilogue::collective::CollectiveBuilder<ArchTag, OpClass, Shape<_256,_128,_32>, Shape<_1,_1,_1>, EpilogueTileAuto, ElementAccumulator, ElementCompute, ElementC, LayoutC, AlignmentC, ElementD, LayoutD, AlignmentD, EpilogueScheduleAuto>::CollectiveOp。
  • Kernel：GemmKernel = gemm::kernel::GemmUniversal<Shape<int,int,int,int>, CollectiveMainloop, CollectiveEpilogue>，Gemm = GemmUniversalAdapter<GemmKernel>。
  • 数据：DeviceAllocation<ElementA/B/C/D>，reset(M*K*L) 等；stride：make_mute_packed_stride(Stride{}, make_shape(M, K, L))；调用：同上（Arguments → workspace → can_implement → initialize → run）。

源码对照（节选自 collective_builder.mu）：

// 类型与布局
using LayoutA = mutlass::layout::ColumnMajor;
using LayoutB = mutlass::layout::RowMajor;
using LayoutC = mutlass::layout::ColumnMajor;
using LayoutD = mutlass::layout::ColumnMajor;
using ElementA = mutlass::half_t;
using ElementB = mutlass::half_t;
using ElementC = mutlass::half_t;
using ElementD = mutlass::half_t;
using ElementAccumulator = float;
using ArchTag = mutlass::arch::Mp22;
using OpClass = mutlass::arch::OpClassTensorOp;
static constexpr int AlignmentA = 16 / sizeof(ElementA);
// ... AlignmentB/C/D 同理

// Collective Builder 构造 Mainloop / Epilogue
using CollectiveMainloop = typename mutlass::gemm::collective::CollectiveBuilder<
    ArchTag, OpClass,
    ElementA, LayoutA, AlignmentA, ElementB, LayoutB, AlignmentB,
    ElementAccumulator,
    Shape<_256, _128, _32>, Shape<_1, _1, _1>,
    mutlass::gemm::collective::StageCountAuto,
    mutlass::gemm::collective::KernelScheduleAuto
>::CollectiveOp;
using CollectiveEpilogue = typename mutlass::epilogue::collective::CollectiveBuilder<
    ArchTag, OpClass,
    Shape<_256, _128, _32>, Shape<_1, _1, _1>,
    mutlass::epilogue::collective::EpilogueTileAuto,
    ElementAccumulator, ElementCompute,
    ElementC, LayoutC, AlignmentC, ElementD, LayoutD, AlignmentD,
    mutlass::epilogue::collective::EpilogueScheduleAuto
>::CollectiveOp;

using GemmKernel = mutlass::gemm::kernel::GemmUniversal<
    Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue
>;
using Gemm = mutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;

// 数据与调用
mutlass::DeviceAllocation<typename Gemm::ElementA> block_A;
block_A.reset(M * K * L);
StrideA stride_A = mutlass::make_mute_packed_stride(StrideA{}, mute::make_shape(M, K, L));
typename Gemm::Arguments arguments{
  mutlass::gemm::GemmUniversalMode::kGemm, problem_size,
  {block_A.get(), stride_A, block_B.get(), stride_B},
  {{options.alpha, options.beta}, block_C.get(), stride_C, block_D.get(), stride_D},
  mutlass::KernelHardwareInfo{}
};
mutlass::device_memory::allocation<uint8_t> workspace(Gemm::get_workspace_size(arguments));
MUTLASS_CHECK(gemm.can_implement(arguments));
MUTLASS_CHECK(gemm.initialize(arguments, workspace.get()));
MUTLASS_CHECK(gemm.run());

示例 3：MP31 FP8 GEMM（02_mp31_fp8_gemm_with_collective_builder）​

在 MP31 上使用 FP8 (E4M3) 与 Collective Builder，带 TME、固定 Stage 与 Epilogue 调度。源码：examples/02_mp31_fp8_gemm_with_collective_builder/mp31_fp8_gemm.mu。

要点：

  • 类型：ElementA/B = float_e4m3_t，ElementC/D = half_t，ElementAccumulator = float；对齐：AlignmentA/B/C = 128 / sizeof_bits<Element>::value。
  • Tile/Cluster：TileShape = Shape<_256,_256,_64>，ClusterShape = Shape<_1,_1,_1>；StageCount：StageCount<4>；KernelSchedule：KernelTmeWarpSpecialized；EpilogueSchedule：epilogue::NoSmem。
  • Epilogue 融合：ThreadEpilogueOp = epilogue::fusion::LinearCombination<ElementC, ElementCompute, ElementC, ElementCompute>；Epilogue 的 CollectiveBuilder 最后一参为 CollectiveMainloop（与 Mainloop 一致）。
  • Kernel：GemmUniversal<Shape<int,int,int>, CollectiveMainloop, CollectiveEpilogue, PersistentScheduler>（ProblemShape 为 3 维表示非 batched）；调用：同上，stride 与 Arguments 构造方式同 01。

源码对照（节选自 mp31_fp8_gemm.mu）：

// A/B 为 FP8 E4M3，C/D 为 half_t
using ElementA = mutlass::float_e4m3_t;
using ElementB = mutlass::float_e4m3_t;
using ElementC = mutlass::half_t;
constexpr int AlignmentA = 128 / mutlass::sizeof_bits<ElementA>::value;
using TileShape = Shape<_256,_256,_64>;
using ClusterShape = Shape<_1,_1,_1>;
using StageCountType = mutlass::gemm::collective::StageCount<4>;
using KernelSchedule = mutlass::gemm::KernelTmeWarpSpecialized;
using EpilogueSchedule = mutlass::epilogue::NoSmem;
using ThreadEpilogueOp = mutlass::epilogue::fusion::LinearCombination<
    ElementC, ElementCompute, ElementC, ElementCompute>;

using CollectiveMainloop = typename mutlass::gemm::collective::CollectiveBuilder<
    ArchTag, OperatorClass,
    ElementA, LayoutA, AlignmentA, ElementB, LayoutB, AlignmentB,
    ElementAccumulator, TileShape, ClusterShape,
    StageCountType, KernelSchedule
>::CollectiveOp;
using CollectiveEpilogue = typename mutlass::epilogue::collective::CollectiveBuilder<
    ArchTag, OperatorClass,
    TileShape, ClusterShape,
    mutlass::epilogue::collective::EpilogueTileAuto,
    ElementAccumulator, ElementAccumulator,
    ElementC, LayoutC, AlignmentC, ElementC, LayoutC, AlignmentC,
    EpilogueSchedule, ThreadEpilogueOp,
    CollectiveMainloop   // 最后一参与 Mainloop 一致
>::CollectiveOp;

using GemmKernel = mutlass::gemm::kernel::GemmUniversal<
    Shape<int,int,int>, CollectiveMainloop, CollectiveEpilogue,
    mutlass::gemm::PersistentScheduler
>;
using Gemm = mutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;

// 非 batched：ProblemShape 为 3 维，stride 与 (M,N,K) 对应
typename Gemm::Arguments arguments{
  mutlass::gemm::GemmUniversalMode::kGemm,
  {options.m, options.n, options.k},
  {block_A.get(), stride_A, block_B.get(), stride_B},
  {{options.alpha, options.beta}, block_C.get(), stride_C, block_D.get(), stride_D}
};
arguments.scheduler.raster_order = options.raster;
arguments.scheduler.swizzle_size = options.swizzle;

示例 4：MP31 FP8 缩放 GEMM（03_mp31_fp8_scaling_gemm）​

在 MP31 上做 FP8 + per-group 缩放 的 GEMM（scale 按 M/N/K 分组）。源码：examples/03_mp31_fp8_scaling_gemm/fp8_scaling_gemm.mu。

要点：

  • 类型：A/B 为 float_e4m3_t，C/D 为 float；Tile：Shape<_128,_256,_128>；StageCount：StageCount<2>。
  • 缩放配置：GroupScaleConfig<ScaleGranularityM, ScaleGranularityN, ScaleGranularityK>，与 Tile 维度需满足整除关系。
  • Kernel：KernelTmeWarpSpecializedScaledAccum<ScaleGranularityM, ScaleGranularityN, ScaleGranularityK>；Epilogue 为 LinearCombination<ElementD, ElementCompute, ElementC, ElementCompute>。
  • 参考实现：主机侧参考在 reference/host/gemm_with_groupwise_scaling.hpp，用于正确性校验。

源码对照（节选自 fp8_scaling_gemm.mu）：

// A/B 为 FP8 E4M3，C/D 为 float
using ElementA = mutlass::float_e4m3_t;
using ElementB = mutlass::float_e4m3_t;
using ElementC = float;
using ElementD = ElementC;
using TileShapeMNK = Shape<_128,_256,_128>;
using StageCountType = mutlass::gemm::collective::StageCount<2>;

// 按 M/N/K 分组的缩放配置，需与 Tile 维度整除
template <int ScaleGranularityM_, int ScaleGranularityN_, int ScaleGranularityK_>
struct GroupScaleConfig {
  using TileShape = TileShapeMNK;
  using KernelSchedule = mutlass::gemm::KernelTmeWarpSpecializedScaledAccum<
      ScaleGranularityM_, ScaleGranularityN_, ScaleGranularityK_>;
  using EpilogueSchedule = mutlass::epilogue::NoSmem;
  using ThreadEpilogueOp = mutlass::epilogue::fusion::LinearCombination<
      ElementD, ElementCompute, ElementC, ElementCompute>;
  // ScaleMsPerTile = size<0>(TileShape{}) / ScaleGranularityM; 等
};
using GroupScale1x128x128Config = GroupScaleConfig<1, 128, 128>;

// Collective Builder + ScaledAccum Kernel
template <class ScheduleConfig>
struct GroupScaleGemm {
  using CollectiveMainloop = typename mutlass::gemm::collective::CollectiveBuilder<
    ArchTag, OperatorClass,
    ElementA, LayoutA, AlignmentA, ElementB, LayoutB, AlignmentB,
    ElementAccumulator, TileShape, ClusterShape,
    StageCountType, typename ScheduleConfig::KernelSchedule
  >::CollectiveOp;
  using CollectiveEpilogue = typename mutlass::epilogue::collective::CollectiveBuilder<
    ArchTag, OperatorClass, TileShape, ClusterShape,
    mutlass::epilogue::collective::EpilogueTileAuto,
    ElementAccumulator, ElementAccumulator,
    ElementC, LayoutC, AlignmentC, ElementD, LayoutD, AlignmentD,
    EpilogueSchedule, ThreadEpilogueOp, CollectiveMainloop
  >::CollectiveOp;
  using GemmKernel = mutlass::gemm::kernel::GemmUniversal<
    Shape<int,int,int>, CollectiveMainloop, CollectiveEpilogue,
    mutlass::gemm::PersistentScheduler
  >;
  using Gemm = mutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
};
using GroupScale1x128x128Gemm = GroupScaleGemm<GroupScale1x128x128Config>;
// 参考：reference/host/gemm_with_groupwise_scaling.hpp

API 参考​

以下与 include/mutlass/、tools/util/include/mutlass/、tools/library/ 中定义一致，仅列出常用入口；完整签名以头文件为准。

核心命名空间与头文件​

命名空间 / 头文件说明
mutlass根命名空间
mutlass::archarch::Mp22、arch::Mp31、arch::OpClassTensorOp 等（include/mutlass/arch/）
mutlass::layoutlayout::RowMajor、layout::ColumnMajor（include/mutlass/detail/layout.hpp）
mutlass::gemmGEMM 入口、枚举（include/mutlass/gemm/）
mutlass::gemm::kernelGemmUniversal<ProblemShape, CollectiveMainloop, CollectiveEpilogue, ...>（include/mutlass/gemm/kernel/）
mutlass::gemm::deviceGemmUniversalAdapter<GemmKernel>（include/mutlass/gemm/device/gemm_universal_adapter.h）
mutlass::gemm::collectiveCollectiveBuilder、CollectiveMma、StageCount 等（include/mutlass/gemm/collective/）
mutlass::epilogue::collectiveEpilogue CollectiveBuilder、EpilogueTileAuto 等（include/mutlass/epilogue/collective/）
mutlass::epilogue::threadLinearCombination（include/mutlass/epilogue/thread/linear_combination.h）
mutlass::epilogue::fusionLinearCombination<ElementOutput, ElementCompute, ElementSource, ElementScalar>（include/mutlass/epilogue/fusion/）
mutlass::util工具类（tools/util/include/mutlass/util/）
mutlass::reference::device设备侧参考：gett、BlockCompareEqual、BlockFillRandomUniform 等

主要类与方法​

DeviceAllocation（mutlass/util/device_memory.h）​

template<typename T>
class DeviceAllocation {
public:
    DeviceAllocation();
    DeviceAllocation(size_t capacity);
    T* get() const;
    size_t size() const;
    void reset();
    void reset(size_t capacity);
    void reset(T* ptr, size_t capacity);
    // copy 等见头文件；设备间拷贝可用 device_memory::copy_device_to_device
};

HostTensor（mutlass/util/host_tensor.h）​

template<typename Element, typename Layout>
class HostTensor {
public:
    HostTensor();
    HostTensor(TensorCoord const& extent, bool device_backed = true);
    HostTensor(TensorCoord const& extent, Layout const& layout, bool device_backed = true);
    Element* data();
    void reset(TensorCoord const& extent, Layout const& layout, bool device_backed = true);
    // extent、layout、device 指针等见头文件
};

GemmUniversalAdapter（include/mutlass/gemm/device/gemm_universal_adapter.h）​

template<class GemmKernel_>
class GemmUniversalAdapter</* 3.x kernel 特化 */> {
public:
    using GemmKernel = GemmKernel_;
    using Arguments = typename GemmKernel::Arguments;
    using Params = typename GemmKernel::Params;
    using ElementA/B/C/D = ...;

    bool can_implement(Arguments const& args) const;
    static size_t get_workspace_size(Arguments const& args);
    Status initialize(Arguments const& args, void* workspace = nullptr, ...);
    Status run();  // 使用已 initialize 的 params_
    Status run(Arguments const& args, void* workspace = nullptr, ...);
    Status operator()(Arguments const& args, void* workspace = nullptr, ...);
};

Arguments 由 Kernel 定义，通常包含：GemmUniversalMode mode、ProblemShape problem_size、mainloop 的 A/B 指针与 stride、epilogue 的 alpha/beta 与 C/D 指针及 stride、KernelHardwareInfo。具体见各 gemm::kernel::GemmUniversal 特化。

工具函数​

  • stride：mutlass::make_mute_packed_stride(Stride{}, mute::make_shape(M, K, L))（util/packed_stride.hpp）。
  • workspace：mutlass::device_memory::allocation<uint8_t> workspace(size)（util/device_memory.h）。
  • 参考 GEMM：mutlass::reference::device::gett(problem_size, A, stride_A, B, stride_B, ..., C, stride_C, D, stride_D, alpha, beta)（util/reference/device/gett.hpp）。

更新日志​

  • 2026-01-29: 创建

版权所有 2024 摩尔线程，本文档受国际版权法保护。

摩尔线程和摩尔线程徽标是摩尔线程智能科技 (北京)有限责任公司的注册商标。

免责声明

本文档提供有关摩尔线程产品的信息。本文档并未授权任何知识产权的许可，并未以明示或暗示，或以禁止反言或其他方式授予任何知识产权许可。