为何小矩阵场景下cuBLAS批量矩阵乘转置性能大幅下降？

我在测试cuBLAS的strided batch矩阵乘法时发现，小矩阵场景下启用转置标志会导致性能大幅下降。使用Julia的CUDA.jl编写了三组对比实验方案：

• 无转置的标准gemm（gemm_strided_batched('N', 'N')）
• 对矩阵B启用转置标志的gemm（gemm_strided_batched('N', 'T')）
• 自定义CUDA核手动转置B后，再执行无转置gemm

测试代码 #

using BenchmarkTools
using CUDA

N = 1000
D = 3

A = CUDA.randn(D, D, N)
B = CUDA.randn(D, D, N)

# Simple kernel to transpose the first two dimensions of a 3D array
function batch_transpose!(B, B_T)
    index = (blockIdx().x - 1) * blockDim().x + threadIdx().x
    stride = blockDim().x * gridDim().x
    for i in index:stride:size(B, 3)
        for j in 1:size(B, 2)
            for k in 1:size(B, 1)
                B_T[k, j, i] = B[j, k, i]
            end
        end
    end
end

# Compute A * B_T by first transposing B explicitly
function transpose_first(A, B)
    B_T = CuArray{Float32}(undef, size(B, 2), size(B, 1), size(B, 3))
    @cuda threads = 256 blocks = ceil(Int, size(B, 3) / 256) batch_transpose!(B, B_T)
    return CUDA.CUBLAS.gemm_strided_batched('N', 'N', A, B_T)
end

# Benchmark the three approaches
bench1 = @benchmark CUDA.@sync CUDA.CUBLAS.gemm_strided_batched('N', 'N', $A,$B)
bench2 = @benchmark CUDA.@sync CUDA.CUBLAS.gemm_strided_batched('N', 'T', $A, $B)
bench3 = @benchmark CUDA.@sync transpose_first($A, $B)

# Print the results
println("Non-transposed gemm: $(median(bench1.times) / 1e6)")
println("Transposed gemm: $(median(bench2.times) / 1e6)")
println("Transposed manually: $(median(bench3.times) / 1e6)")

测试环境与结果 #

测试在RTX 4090显卡上进行，两组不同矩阵尺寸的测试结果如下：

D=3（小矩阵） #

Non-transposed gemm: 9.854 μs
Transposed gemm    : 25.248 μs
Transposed manually: 14.011 μs

启用转置标志的gemm性能远低于无转置版本，甚至不如手动转置后执行gemm的方案。

D=30（中等矩阵） #

Non-transposed gemm: 38.468 μs
Transposed gemm    : 38.357 μs
Transposed manually: 309.046 μs

此时无转置与转置gemm的性能差距消失，手动转置方案性能则显著下降。

问题 #

对批量乘转置这种常见操作在cuBLAS中小矩阵场景下的糟糕表现感到惊讶，想了解背后的原因。

内容的提问来源于stack exchange，提问作者Tim Hargreaves