跳转至

Performance Tips

1. Basic Concepts

1.1 Cooperative Vector API

这是一种编程模型，允许每个线程独立执行 矩阵-向量乘法（Matrix-Vector Multiplication, M-V）。每个线程负责一个输入向量，通过协作完成多个向量的计算。

1.2 Tensor Cores

功能：Tensor Cores 是 NVIDIA GPU 中专门用于加速矩阵运算的硬件单元，特别擅长执行 矩阵-矩阵乘法（Matrix-Matrix multiplication）。

并行粒度：一个完整的 wave（波前）或 warp（线程束） 可以被用来一次性加载和处理整个矩阵块，从而实现高吞吐量。

虽然传统的协同向量 API（CoopVec）允许每个线程独立执行矩阵-向量乘法，但其效率较低；而 Tensor Cores 更擅长处理高吞吐的矩阵-矩阵乘法，尤其适用于低精度计算（如 FP16、FP8）。

2. GPU 执行模型与线程发散

NVIDIA GPU 将执行线程组织成 32 个线程组成的"warp"，这是 SIMD/SIMT 架构的基本调度单位。每个 warp 在一个时钟周期内执行一条指令，但作用于不同的数据（即 SIMT 模型）。

所有线程在同一个 wave 中执行同一条指令

  • 在一个 wave 内的所有线程 同时执行相同的指令，但操作的是各自的数据。
  • 例如：add(a[i], b[i]) 中，每个线程处理不同索引 i 的数据。
  • 有些线程可能因条件判断而被标记为 inactive（非活跃），不参与当前计算，但仍占用资源。

2.1 执行发散（Execution Divergence）

当一个 warp 中的线程进入不同的代码分支时，就会发生执行发散。这会导致分支被串行化执行，部分线程处于空闲状态，严重影响性能。

if (threadIdx.x < 4) {
    A;
    B;
} else {
    X;
    Y;
}
Z;

这里前 4 个线程执行 A; B;，其余执行 X; Y;。

因为 GPU 必须按顺序执行，所以这两个分支会被串行化执行（serialized）：

  • 先执行 A; B;（前 4 个线程），其他线程等待；
  • 再执行 X; Y;（后 28 个线程），前 4 个线程等待；
  • 最后所有线程一起执行 Z;。

执行效率下降，因为部分线程处于空闲状态。

2.2 数据发散（Data Divergence）

数据发散涉及内存访问模式。最优内存访问发生在多个线程访问连续且对齐的内存地址时，这样可以将整个 warp 的内存请求合并为少量内存事务，极大提升带宽利用率。

如果线程访问的地址不规则或未对齐，则会导致数据发散（data divergence）。

3. Tensor Core 优化策略

Tensor Cores 以整个 wave 为单位执行单个矩阵乘累加（MMA）操作，因此要求所有线程使用相同的矩阵输入。

3.1 避免矩阵发散

当使用 Cooperative Vector API 时，若不同线程的矩阵输入发散，会导致驱动程序将其操作串行化，严重降低性能。

建议做法：

  • 按材质分组绘制调用：确保相同材质的几何体一起处理，共享相同矩阵
  • 手动排序线程或使用 Shader Execution Reordering（SER）：通过重新排列线程执行顺序，使具有相同矩阵的线程聚集在一起

3.2 矩阵布局优化

Tensor Cores 并不直接支持标准的行主序或列主序矩阵存储方式。在执行 MMA 操作之前，矩阵元素必须被重新排列以满足硬件要求。

  • 权重矩阵必须在写入内存前完成预重排
  • 输入和输出数据可通过现代图形 API 自动处理 shuffle 操作
  • 一个 MMA 的输出可直接作为另一个 MMA 的输入，无需额外 shuffle

开发者可以将标准矩阵上传至 GPU，然后通过 CoopVec API 转换为适合 Tensor Core 的优化布局：

cCoopVec<half, OutputSize> eval(CoopVec<half, InputSize> input)
{
    let output = coopVecMatMulAdd<half, OutputSize>(
        input, CoopVecComponentType.Float16,
        weights, CoopVecComponentType.Float16,
        biases, CoopVecComponentType.Float16,
        CoopVecMatrixLayout.InferencingOptimal,
        false, // is matrix transposed
        sizeof(half) * InputSize); // matrix stride
    return max(output * 0.01h, output);
}

4. Mapping MLP onto Tensor cores

将多层感知机(MLP)映射到 Tensor Cores 需要在 MMA 操作前后添加 shuffle 操作：

cshflInput = shuffle(input);
coopVecMatMul(temp0, shflInput, weightBuffer, offset0);
layer0 = unshuffle(temp0);
layer0 = max(layer0, 0);

shflLayer0 = shuffle(layer0);
coopVecMatMul(temp1, shflLayer0, weightBuffer, offset1);
layer1 = unshuffle(temp1);
layer1 = max(layer1, 0);

shflLayer1 = shuffle(layer1);
coopVecMatMul(temp2, shflLayer1, weightBuffer, offset2);
output = unshuffle(temp2);

这是一个 Base 版本，因为每次 shuffle 都会引入额外的指令、带宽消耗和延迟。如果不优化，这些冗余操作会显著降低整体性能。

由于 max() 是逐元素操作，可以在 shuffled 格式下直接执行，无需还原：

cshflinput = shuffle(input);
coopVecMatMul(temp0, shflinput, weightBuff, offset0);
temp0_1 = max(temp0, 0);
coopVecMatMul(temp1, temp0_1, weightBuff, offset1);
temp1_1 = max(temp1, 0);
coopVecMatMul(temp2, temp1_1, weightBuff, offset1);
output = unshuffle(temp2);

层融合（Layer Fusion）

将相邻的矩阵乘法和激活函数等操作合并成一个整体流程，并尽可能保持数据在 Tensor Core 所需的布局中传递，避免不必要的格式转换。

层融合有时并不可行

虽然层融合能带来巨大性能提升，但它 并不总是可行。失败原因包括：

  • 某些激活函数或操作无法在 shuffled 布局下正确执行；
  • 需要将中间结果输出到外部（如调试、可视化）；
  • 不同层的权重布局不兼容；

此版本假设所有线程使用相同的权重偏移（即同一组参数），适用于 同质化输入 场景。

当存在 多个不同的权重偏移（offsets） 时（例如：不同材质、不同实例、不同模型变体），每个线程可能需要访问不同的权重块。

如果直接运行上述代码，会导致：

  • 不同线程使用不同的 offset；
  • 引发 矩阵发散（divergent matrices）；
  • 驱动程序会将这些操作串行化处理（如之前章节所述），严重降低性能。

当存在多个不同的权重偏移时，需要对每种唯一组合单独执行完整的 MLP 计算流程：

cshflinput = shuffle(input);
foreach (unique combination of offsets)
{
    coopVecMatMul(temp0, shflinput, weightBuff, offset0);
    temp0_1 = max(temp0, 0);

    coopVecMatMul(temp1, temp0_1, weightBuff, offset1);
    temp1_1 = max(temp1, 0);

    coopVecMatMul(temp2, temp1_1, weightBuff, offset2);
    mergeThreadResults(temp2);
}
output = unshuffle(temp2);

5. Preventing Layer Fusion Failures

在使用 Cooperative Vector（CoopVec）API 进行神经网络推理时，因避免中间操作不当而导致 层融合（Layer Fusion）失败。

案例1：

coopVecMatMul(layer0, input, weightBuffer, offset0);

for (int i = 0; i < 64; ++i)
{
    // Leaky ReLU
    if (layer0[i] < 0)
        layer0[i] *= 0.01;
}

coopVecMatMul(layer1, layer0, weightBuffer, offset1);

使用了传统的 C-style 循环遍历每个元素。每次访问 layer0[i] 都是独立的内存读写；线程可能发散（divergence），因为不同线程的条件判断结果不同；

一个解决方法就是使用向量内建函数

coopVecMatMul(layer0, input, weightBuffer, offset0);

// Use vector operations to express the same math
layer0 = max(layer0, layer0 * 0.01);

coopVecMatMul(layer1, layer0, weightBuffer, offset1);

有用的向量内建函数

min, max, clamp, step 等都是支持向量化的；

可以在 hlsl.meta.slang 文件中找到完整列表；

案例2：

coopVecMatMul(layer0, input, weightBuffer, offset0);

CoopVec<half, 64> bias;
for (int i = 0; i < 64; ++i)
{
    bias[i] = biasBuffer.load<half>(biasOffset0 + i * sizeof(half));
}

layer0 += bias;
layer0 = max(layer0, 0); // ReLU

coopVecMatMul(layer1, layer0, weightBuffer, offset1);

上述代码从全局缓冲区中加载偏置（bias）值，并逐个赋值给 bias[i]。

  • 使用了 逐元素循环（element-wise loop）；
  • 每次调用 load<half>() 都是独立的内存事务；
  • 编译器无法识别这是一个批量加载操作；

进而导致层融合失败，系统被迫在 coopVecMatMul 之间插入 unshuffle 和 shuffle 操作，性能下降。

一个解决方法是使用向量加载（Vector Load）

coopVecMatMul(layer0, input, weightBuffer, offset0);

let bias = CoopVec<half, 64>.load(biasBuffer, biasOffset0);
layer0 += bias;

layer0 = max(layer0, 0); // ReLU

coopVecMatMul(layer1, layer0, weightBuffer, offset1);
  • 使用 CoopVec<half, 64>.load() 一次性加载整个偏置向量；
  • 所有 64 个元素被并行读取，形成一次高效的内存事务；
  • 数据保持在 shuffled 布局下；

这样整体流程连续，无中断；驱动能够识别这是一个可融合的操作，允许后续 coopVecMatMul 继续使用相同布局

上述的代码仍然有改进的空间，我们可以使用 coopVecMatMulAdd 函数，它将 矩阵乘法 + 偏置加法 合并为一个单一指令。

coopVecMatMulAdd(layer0, input, weightBuffer, offset0,
                 biasBuffer, biasOffset0);

layer0 = max(layer0, 0); // ReLU

coopVecMatMul(layer1, layer0, weightBuffer, offset1);

评论区

对你有帮助的话请给我个赞和 star => GitHub stars
欢迎跟我探讨！！！