矩阵Q,K,V 在HBM,SRAM芯片大小为M,其中N=5,d=3, M=59

1.设置块大小Bc=4,Br=3

2.初始化O,l,m

3.分割Q为Tr块[Q0,Q1],每块大小Br*d,分割K,V为Tc块[K0,K1],[V0,V1],每块大小Bc*d

4.分割O为Tr块,每块大小Br * d,分割l为Tr块,每块大小Br,分割m为Tr,每块大小为Br

外层循环:
从HBM加载K0,V0到SRAM

内层循环:
从HDM中加载Q0,O0,l0,m0到SRAM中

在芯片中,计算S[0,0] = Q0 * K0

在芯片中,计算当前块每一行的 最大值m[0,0] = rowmax(S[0, 0]),Softmax分子P[0, 0] = exp(S[0,0] - m[0,0]),Softmax分母l[0, 0]=rowsum(P[0, 0])



在芯片中计算出:l[0, 0], m[0, 0]和l[0,0]的每一行的和l0

计算出O[0],将O[0]结果写入到HBM中


从HBM中将Q1,O1,l1,m1加载到SRAM中






