Multi-Head Attention

原始题目：LeetGPU - Multi-Head Attention

题目描述

实现多头自注意力（Multi-Head Self-Attention）。给定查询矩阵 $Q$、键矩阵 $K$ 和值矩阵 $V$（大小均为 $N \times d_{model}$），计算：

$$\text{MultiHead}(Q, K, V) = \text{Concat}(head_1, \dots, head_h)$$

其中每个头计算：

$$head_i = \text{softmax}\!\left(\frac{Q_i K_i^T}{\sqrt{d_k}}\right) V_i$$

$d_k = d_{model} / h$，$Q_i, K_i, V_i$ 是输入矩阵的第 $i$ 个头分区。

实现要求

• 不允许使用外部库。
• solve 函数签名必须保持不变。

示例

N=2, d_model=4, h=2
Q=[[1,0,2,3],[4,5,6,7]], K=[[1,2,3,4],[5,6,7,8]], V=[[0.5,1,1.5,2],[2.5,3,3.5,4]]
Output (2×4): [[2.39,2.89,3.50,4.00],[2.50,3.00,3.50,4.00]]

约束条件

• $1 \le N \le 10{,}000$，$2 \le d_{model} \le 1{,}024$，$1 \le h \le d_{model}$，$d_{model} \bmod h = 0$。
• 性能测试在 $N=1{,}024, d_{model}=1{,}024$ 下进行。

解题思路

MHA 是 Transformer 的核心算子。每个头独立计算 softmax attention，最后拼接输出。GPU 上高效实现的关键是将多个小矩阵乘法 batch 化：将 $Q, K, V$ 的重塑视为 batched GEMM。每个头的数据量较小（$d_k$ 一般 64–128），可以多个头并行处理。FlashAttention 通过分块计算 $QK^T$ 避免完整 attention 矩阵的显存分配，是生产级实现的标准方案。

代码实现

CUDA

#include <cuda_runtime.h>
#include <math.h>
__global__ void mha_kernel(const float* Q, const float* K, const float* V, float* O,
    int N, int d_model, int h) {
    int dk = d_model / h;
    int head = blockIdx.x, i = threadIdx.x;
    if(head >= h || i >= N) return;
    float scale = 1.0f / sqrtf((float)dk);
    // Each head: softmax(Q_i @ K_i^T / sqrt(dk)) @ V_i
    float scores[256], mx = -INFINITY, sm = 0.0f;
    for(int j = 0; j < N; j++) {
        float dot = 0.0f;
        for(int k = 0; k < dk; k++)
            dot += Q[(i*d_model) + head*dk + k] * K[(j*d_model) + head*dk + k];
        scores[j] = dot * scale;
        mx = fmaxf(mx, scores[j]);
    }
    for(int j = 0; j < N; j++) sm += expf(scores[j] - mx);
    for(int k = 0; k < dk; k++) {
        float sum = 0.0f;
        for(int j = 0; j < N; j++)
            sum += expf(scores[j] - mx) / sm * V[(j*d_model) + head*dk + k];
        O[(i*d_model) + head*dk + k] = sum;
    }
}
extern "C" void solve(const float* Q, const float* K, const float* V, float* O, int N, int d_model, int h) {
    mha_kernel<<<h, N>>>(Q, K, V, O, N, d_model, h);
    cudaDeviceSynchronize();
}

Triton

import triton, triton.language as tl
@triton.jit
def mha_kernel(Q_ptr,K_ptr,V_ptr,O_ptr, N,d_model,h, BLOCK:tl.constexpr):
    head=tl.program_id(0); pos=tl.program_id(1)*BLOCK+tl.arange(0,BLOCK)
    mask=pos<N; dk=d_model//h; d_range=tl.arange(0,dk)
    q=tl.load(Q_ptr+pos[:,None]*d_model+head*dk+d_range[None,:],mask=mask[:,None])
    scale=1.0/tl.sqrt(dk.to(tl.float32)); acc=tl.zeros((BLOCK,dk),tl.float32)
    for j in range(N):
        k=tl.load(K_ptr+j*d_model+head*dk+d_range); v=tl.load(V_ptr+j*d_model+head*dk+d_range)
        score=tl.sum(q*k[None,:],axis=1)*scale; w=tl.softmax(score)
        acc+=w[:,None]*v[None,:]
    tl.store(O_ptr+pos[:,None]*d_model+head*dk+d_range[None,:],acc,mask=mask[:,None])