Batched Matrix Multiplication

agicy2026/6/6大约 1 分钟

Batched Matrix Multiplication

原始题目：LeetGPU - Batched Matrix Multiplication

题目描述

在 FP32 中实现批量矩阵乘法。给定一批形状为 $[B, M, K]$ 的矩阵 $A$ 和一批形状为 $[B, K, N]$ 的矩阵 $B$ ，计算输出批次 $C$ （形状 $[B, M, N]$ ），使得对每个批次索引 $b$ ：

C_b = A_b \times B_b

所有矩阵以行优先顺序存储，使用 32 位浮点数（FP32）。

实现要求

不允许使用外部库。
solve 函数签名必须保持不变。
最终结果必须存储在数组 $C$ 中。

示例

Input:  B=2, M=2, K=3, N=2
        A[0] = [[1,2,3],[4,5,6]], A[1] = [[7,8,9],[10,11,12]]
        B[0] = [[1,2],[3,4],[5,6]], B[1] = [[6,5],[4,3],[2,1]]
Output: C[0] = [[22,28],[49,64]], C[1] = [[92,68],[128,95]]

约束条件

$1 \le B \le 128$ 。
$1 \le M, N, K \le 1{,}024$ 。
性能测试在 $K = M = N = 256$ 的规模下进行。

批量矩阵乘法是对 $B$ 个独立矩阵对同时执行 GEMM。最简单的做法是串行处理每个 batch，但 GPU 的优势在于并行——可以让每个 batch 由不同的 SM 处理，或者使用 cuBLAS 的 cublasGemmBatched / cublasGemmStridedBatched。手写实现中，关键是将 batch 索引作为第三维扩展到已有的二维分块 GEMM 策略中，确保所有 batch 同时利用 GPU 的并行性。

代码实现

CUDA

#include <cuda_runtime.h>
__global__ void batched_matmul(const float* A, const float* B, float* C, int B, int M, int K, int N) {
    int b = blockIdx.z, row = blockIdx.y*blockDim.y+threadIdx.y, col = blockIdx.x*blockDim.x+threadIdx.x;
    if(b < ::B && row < M && col < N) {
        float sum=0.0f;
        for(int i=0;i<K;i++) sum+=A[b*M*K+row*K+i]*B[b*K*N+i*N+col];
        C[b*M*N+row*N+col]=sum;
    }
}
extern "C" void solve(const float* A, const float* B, float* C, int B, int M, int K, int N) {
    dim3 t(16,16), b_((N+15)/16,(M+15)/16,B);
    batched_matmul<<<b_,t>>>(A,B,C,B,M,K,N); cudaDeviceSynchronize();
}

Triton

import triton, triton.language as tl
@triton.jit
def batched_matmul(A_ptr,B_ptr,C_ptr, B,M,K,N, BLOCK_M:tl.constexpr,BLOCK_N:tl.constexpr,BLOCK_K:tl.constexpr):
    pb=tl.program_id(2);pm=tl.program_id(0);pn=tl.program_id(1)
    rm=pm*BLOCK_M+tl.arange(0,BLOCK_M);rn=pn*BLOCK_N+tl.arange(0,BLOCK_N);rk=tl.arange(0,BLOCK_K)
    acc=tl.zeros((BLOCK_M,BLOCK_N),tl.float32)
    for k in range(0,K,BLOCK_K):
        a=tl.load(A_ptr+pb*M*K+rm[:,None]*K+(k+rk)[None,:])
        b=tl.load(B_ptr+pb*K*N+(k+rk)[:,None]*N+rn[None,:])
        acc+=tl.dot(a,b)
    tl.store(C_ptr+pb*M*N+rm[:,None]*N+rn[None,:],acc)