Grouped Query Attention

原始题目：LeetGPU - Grouped Query Attention

题目描述

实现分组查询注意力（GQA），这是 LLaMA-3、Mistral、Gemma 等现代大语言模型使用的注意力机制。GQA 通过在查询头组之间共享键和值头来减少推理时的 KV-cache 内存占用。

给定查询张量 $Q$（$num\_q\_heads$ 个头）和键/值张量 $K$、$V$（各 $num\_kv\_heads$ 个头），计算缩放点积注意力，其中每 $num\_q\_heads / num\_kv\_heads$ 个连续查询头共享同一个键值头：

$$\begin{aligned} scale &= 1 / \sqrt{head\_dim} \\ scores &= Q @ K^T \cdot scale \\ weights &= \text{softmax}(scores) \\ output &= weights @ V \end{aligned}$$

所有张量使用 float32。

实现要求

• 实现 solve(Q, K, V, output, num_q_heads, num_kv_heads, seq_len, head_dim)。
• 不允许使用外部库。
• $num\_q\_heads$ 始终可被 $num\_kv\_heads$ 整除。

示例

num_q_heads=4, num_kv_heads=2 (每组2), seq_len=3, head_dim=4
Q0: [[1,0,0,1],[0,1,1,0],[1,1,0,0]]  →  Q0,Q1 attend to K0,V0
Q1: [[0,1,0,1],[1,0,1,0],[0,0,1,1]]  →  Q2,Q3 attend to K1,V1
...
Output: (4 heads × 3 positions × 4 dims, 结果取2位小数)

约束条件

• $1 \le num\_kv\_heads \le num\_q\_heads \le 64$。
• $1 \le seq\_len \le 4{,}096$，$8 \le head\_dim \le 256$（8 的倍数）。
• 性能测试在 $num\_q\_heads=32$，$num\_kv\_heads=8$，$seq\_len=1{,}024$，$head\_dim=128$ 下进行。

解题思路

GQA 的关键区别在于 KV 头的广播——每个 KV 头被多个 Q 头共享。实现中可以在 batch 维度上复制 KV 头以匹配 Q 头数量，然后执行标准的 MHA；更高效的做法是直接在注意力计算中处理分组映射，避免额外的内存拷贝。KV-cache 友好的内存布局是生产级实现的重点。

代码实现

CUDA

#include <cuda_runtime.h>
#include <math.h>
__global__ void gqa_kernel(const float* Q,const float* K,const float* V,float* O,
    int nqh,int nkvh,int S,int d) {
    int h=blockIdx.x,qperkv=nqh/nkvh;
    int i=threadIdx.x; if(i>=S)return;
    int kvh=h/qperkv; float scale=1.0f/sqrtf(d);
    float scores[256],mx=-INFINITY,sm=0.0f;
    for(int j=0;j<S;j++){
        float dot=0; for(int k=0;k<d;k++)dot+=Q[h*S*d+i*d+k]*K[kvh*S*d+j*d+k];
        scores[j]=dot*scale; mx=fmaxf(mx,scores[j]);
    }
    for(int j=0;j<S;j++)sm+=expf(scores[j]-mx);
    for(int k=0;k<d;k++){float sum=0;for(int j=0;j<S;j++)sum+=expf(scores[j]-mx)/sm*V[kvh*S*d+j*d+k];O[h*S*d+i*d+k]=sum;}
}
extern "C" void solve(const float* Q,const float* K,const float* V,float* O,int nqh,int nkvh,int S,int d) {
    gqa_kernel<<<nqh,S>>>(Q,K,V,O,nqh,nkvh,S,d); cudaDeviceSynchronize();
}

Triton

import triton, triton.language as tl
@triton.jit
def gqa_kernel(Q_ptr,K_ptr,V_ptr,O_ptr, nqh,nkvh,S,d, BLOCK:tl.constexpr):
    h=tl.program_id(0); i=tl.program_id(1)*BLOCK+tl.arange(0,BLOCK)
    mask=i<S; qpk=nqh//nkvh; kvh=h//qpk
    d_range=tl.arange(0,d)
    qi=tl.load(Q_ptr+h*S*d+i[:,None]*d+d_range[None,:],mask=mask[:,None])
    scale=1.0/tl.sqrt(d.to(tl.float32)); acc=tl.zeros((BLOCK,d),tl.float32)
    for j in range(S):
        kj=tl.load(K_ptr+kvh*S*d+j*d+d_range); vj=tl.load(V_ptr+kvh*S*d+j*d+d_range)
        score=tl.sum(qi*kj[None,:],axis=1)*scale
        w=tl.softmax(score); acc+=w[:,None]*vj[None,:]
    tl.store(O_ptr+h*S*d+i[:,None]*d+d_range[None,:],acc,mask=mask[:,None])