CS149-Assignment-4(2023FALL)

CS149-Assignment-4(2023FALL)。

CS149：https://gfxcourses.stanford.edu/cs149/fall23/

Assignment 4: NanoGPT149

该Asst的任务是实现和优化一个基于transformer的神经网络模型的核心部分。这个DNN非常小，但是和ChatGPT这样的大模型的基本原理是一样的。具体而言，需要实现模型的注意力层，并专注于优化提升计算强度，减少内存占用，利用CPU的多核和SIMD并行能力。

这个Asst会接触DNN的实现细节，类似于One API或CuDNN的内核实现。并展现循环分块和循环融合等关键局部性优化的价值。

本任务中执行的NanoGPT模块是一个序列到序列的模型。输入是一系列单词，例如短语 “The course of true love never did run smooth”。模型的输出是一个新的单词序列，可能是根据已经在大量莎士比亚文本上进行训练的模型确定的输入后续。例如，给定上述前缀，模型的输出可能是 “whispered cs149 students whilst coding on assignments”。

NanoGPT模型使用一种称为transformer的流行深度神经网络模块，而其关键组件是被称为注意机制的块。在此任务中，任务是实现这个注意机制。从一个简单的顺序注意力实现开始，在任务的过程中，指导书会引导添加一些优化，如循环分块、循环融合和基本并行化。

Background

注意力机制的输入是三个矩阵Q，K和V，分别对应query，key，value向量。每个矩阵都是N×d大小。N是输入序列的token数(单词数)，每一行都是一个向量，是单词的神经编码。

为了提高模型的效率和表达能力，该注意机制通常会并行运行多次，因为存在多个注意头和批处理中的多个输入。确切了解为什么这样做并不重要，在这个简单的实现中，这些矩阵将呈现为 4D 张量，其中只需关注其中的两个维度（对应于N和d）。

注意力模型的第一步是计算单词之间的interactions，这是通过Q和K的相乘完成的：

$$S = QK^T$$

第二步是对S的每一行进行softmax，softmax会产生每一行的归一化概率。假设是对1维向量计算：

$$softmax(x) = \frac {f(x)}{l(x)} \\ f(x) = [e^{x1} \ e^{x2} \ e^{x3} ...\ e^{xN}] \\ l(x) = \sum^N_{i=1} {f(x)_i}$$

通常，为了提高数值稳定性，Softmax 操作的计算中会减去输入向量中的最大值。这样做可以防止指数函数的过度增长，避免数值上溢。不过为了实现的简单，恶意只使用上述的公式。上述的计算会产生一个注意力权重矩阵P。

最终，利用上述的注意力权重汇总出一组学习值向量，形式还是N×d的矩阵。

$$O = PV$$

总之，注意力层由昂贵的矩阵乘法、softmax和另外一个矩阵乘法组成。这三个部分将是实现过程中的主要部分。

Warm-Up: Accessing Tensors

tensor(张量)是python中的一种数据抽象，实际上是多维数组。python将多维数组抽象为张量，在PyTorch时就不用关注访问数值或矩阵乘法等内部问题。并且Python的张量可轻松移植到GPU。在本次的程序中，只使用CPU，因此使用C++ vector，而不是张量。

理解tensor的核心是知道如何访问tensor。在Warm-Up中，需要编写一个4Dtensor的访问器。程序框架包含一个formatTensor函数，将tensor转换为C++Vector，还有函数将输出Vector转换为tensor。

Step 1: Understand a 2D Accessor

2D的数据访问很好理解，数据是以1D的形式存储的。因此访问2D数组，(i，j)的位置就是[i*num_col+j]。

Step 2: Implement a 4D Accessor

在需要实现的LLM模型中，数组是4维的，因此需要实现4D的数据访问。和2D的原理实际上是一样的。首先考虑3D的情况，(x，y，z)的位置应该是[x*SIZE_X*SIZE_Y+y*SIZE_Y+z]，其中的SIZE_X，SIZE_Y对应的是列长度和y方向上的长度，那么4D也是同理。因此访问4D的(x，y，z，b)的方式应该为：

// Step #2: Implement Read/Write Accessors for a 4D Tensor
inline float fourDimRead(std::vector<float> &tensor, int &x, int &y, int &z, int &b, 
        const int &sizeX, const int &sizeY, const int &sizeZ) {
    return tensor[x*sizeX*sizeY*sizeZ+y*sizeY*sizeZ+z*sizeZ+b];
}

inline void fourDimWrite(std::vector<float> &tensor, int &x, int &y, int &z, int &b, 
        const int &sizeX, const int &sizeY, const int &sizeZ, float &val) {
    tensor[x*sizeX*sizeY*sizeZ+y*sizeY*sizeZ+z*sizeZ+b] = val;
    return; 
}

4D张量的布局应该是按照b，z，y，x的顺序一维展开，这样更符合局部性，相邻元素是连续存储的。

Part 1: A Simple (But Not So Efficient) Implementation of Attention

第一步首先要实现没有优化的串行注意力层。框架中给出了如何访问4Dtensor的示例：

//Here is an example of how to read/write 0's to  Q (B, H, N, d) using the 4D accessors
//loop over Batch Size
for (int b = 0; b < B; b++) {
    //loop over Heads
    for (int h = 0; h < H; h++) {
        //loop over Sequence Length
        for (int i = 0; i < N; i++) {
            //loop over Embedding Dimensionality
            for (int j = 0; j < d; j++) {
                float val = fourDimRead(Q, b, h, i, j, H, N, d);
                val = 0.0;
                fourDimWrite(Q, b, h, i, j, H, N, d, val);
            }
        }
    }
}

这里需要注意，正如2D数组访问需要使用的是num_cols，即第二维的SIZE一样，传入4Dtensor访问(b，h，n，d)的SIZE X Y Z应该是H N d，而不需要知道Batch SIZE。

对于需要实现的注意力层，步骤如下：

STEP1

For each Batch: For each Head: 遍历 Q 和 K，并将Q与K^t相乘，将结果存储在 QK^t 中。QK^t已经预先分配，并作为参数传递给 myNaiveAttention。经过Batch和Head索引后得到的是(N, d)的Q和K的2D矩阵。还要注意K的维度为(N, d)，而想要的K^t的维度为(d, N)。可以不进行转置，直接调整矩阵相乘时的行列顺序就可以了。

STEP2

在获得QK^t(N, N)之后，遍历每一行。对于每一行，获取每个行元素的指数。然后，将每个结果指数除以其行中所有指数的总和，然后将其存储回 QK^t。

STEP3

最后，将QK^t与V进行矩阵乘法，并将结果存储到O中。与Q和K一样，索引Batch和Head后，V和O将具有形状(N, d)。在将QK^t(N, N)与V(N, d)相乘后直接存入O。此外O张量会初始化为0。

下面按照上述注意力层步骤一步一步完成注意力层的计算：

torch::Tensor myNaiveAttention(torch::Tensor QTensor, torch::Tensor KTensor, torch::Tensor VTensor, torch::Tensor QK_tTensor,
                int B, int H, int N, int d){

    // Q, K, V are passed in with Shape: (B, H, N, d)
    //QK^t Intermediate Tensor has Shape (N, N)
    
    //Make O Tensor with Shape (B, H, N, d) 
    at::Tensor OTensor = at::zeros({B, H, N, d}, at::kFloat);

    //Format O, Q, K, and V tensors into 4D vectors
    std::vector<float> O = formatTensor(OTensor);
    std::vector<float> Q = formatTensor(QTensor);
    std::vector<float> K = formatTensor(KTensor);
    std::vector<float> V = formatTensor(VTensor);

    //Format QK_t Tensor into a 2D vector.
    std::vector<float> QK_t = formatTensor(QK_tTensor);
    
    // -------- YOUR CODE HERE  -------- //
    //caculate QK_t = Q(N,d) * K^t(d, N)
    //loop over Batch Size
    for(int b = 0; b < B; b++){
    	//loop over Heads
    	for(int h = 0; h < H; h++){
    		//loop over Sequence Length ijk calculate S(i,j) = Q(i,k)*Q(j,k)
    		for(int i = 0; i < N; i++){
				for(int j = 0; j < N; j++){
                    float val = 0;
					for(int k = 0; k < d; k++){
						val += fourDimRead(Q, b, h, i, k, H, N, d) * fourDimRead(K, b, h, j, k, H, N, d);
					}
                    twoDimWrite(QK_t, i, j, N, val);
				}
			}
            //softmax
            for(int i = 0; i < N ; i++){
                float sum = 0.0;
                for(int j = 0; j < N; j++){
                    sum += exp(twoDimRead(QK_t, i, j, N));
                }
                for(int j = 0; j < N; j++){
                float val =  exp(twoDimRead(QK_t, i, j, N))/sum;
                twoDimWrite(QK_t, i, j, N, val);
                }
            }
    		//loop over Sequence Length ikj calculate O(i,j) = QK_t(i,k)*V(k,j)
    		for(int i = 0; i < N; i++){
				for(int k = 0; k < N; k++){
					for(int j = 0; j < d; j++){
                        float val = fourDimRead(O, b, h, i, j, H, N, d);
						val += twoDimRead(QK_t, i, k, N) * fourDimRead(V, b, h, k, j, H, N, d);
                        fourDimWrite(O, b, h, i, j, H, N, d, val);
                    }
				}
			}
    	}
    }    
    // DO NOT EDIT THIS RETURN STATEMENT //
    // It formats your C++ Vector O back into a Tensor of Shape (B, H, N, d) and returns it //
    return torch::from_blob(O.data(), {B, H, N, d}, torch::TensorOptions().dtype(torch::kFloat32)).clone();
}

最后通过python3 gpt149.py part1 -N 1000进行测试，因为循环的顺序选择了更符合局部性的，所以比reference的性能更好，CPU时间小于reference的50%。

REFERENCE - NAIVE ATTENTION statistics
cpu time:  171.365ms
mem usage:  4512000 bytes

STUDENT - NAIVE ATTENTION statistics
cpu time:  83.875ms
mem usage:  4512000 bytes

Part 2: Blocked Matrix Multiply and Unfused Softmax

part2要改进上面的实现，采用分块矩阵乘法。这里getconf LEVEL1_DCACHE_LINESIZE看了一下缓存行是64字节，能存16个浮点，所以分块的行可以取16个元素，每次取64行。

torch::Tensor myUnfusedAttentionBlocked(torch::Tensor QTensor, torch::Tensor KTensor, torch::Tensor VTensor, torch::Tensor QK_tTensor,
                int B, int H, int N, int d){
    
    // Q, K, V are passed in with Shape: (B, H, N, d)
    //QK^t Intermediate Tensor has Shape (N, N)

    //Make O Tensor with Shape (B, H, N, d) 
    at::Tensor OTensor = at::zeros({B, H, N, d}, at::kFloat);

    //Format O, Q, K, and V tensors into 4D vectors
    std::vector<float> O = formatTensor(OTensor);
    std::vector<float> Q = formatTensor(QTensor);
    std::vector<float> K = formatTensor(KTensor);
    std::vector<float> V = formatTensor(VTensor);

    //Format QK_t Tensor into a 2D vector.
    std::vector<float> QK_t = formatTensor(QK_tTensor);

    // -------- YOUR CODE HERE  -------- //
    //caculate QK_t = Q(N,d) * K^t(d, N)
    //loop over Batch Size
    for(int b = 0; b < B; b++){
    	//loop over Heads
    	for(int h = 0; h < H; h++){
    		//loop over Sequence Length ijk calculate S(i,j) = Q(i,k)*Q(j,k)
            std::fill(QK_t.begin(), QK_t.end(), 0);
            int bsize_H, bsize_W;
            bsize_W = 16, bsize_H = 64;
            for(int is = 0; is < N; is += bsize_H){
                for(int js = 0; js < N; js += bsize_W){
                    for(int ks = 0; ks < d; ks += bsize_W){
                        for(int i = is; i < std::min(is+bsize_H, N); i++){
                            for(int j = js; j < std::min(js+bsize_W, N); j++){
                                float val = twoDimRead(QK_t, i, j, N);
                                for(int k = ks; k < std::min(ks+bsize_W, d); k++){
                                    val += fourDimRead(Q, b, h, i, k, H, N, d) * fourDimRead(K, b, h, j, k, H, N, d);
                                }
                                twoDimWrite(QK_t, i, j, N, val);
                            }
                        }
                    }
                }
            }
            //softmax
            for(int i = 0; i < N ; i++){
                float sum = 0.0;
                for(int j = 0; j < N; j++){
                    sum += exp(twoDimRead(QK_t, i, j, N));
                }
                for(int j = 0; j < N; j++){
                float val =  exp(twoDimRead(QK_t, i, j, N))/sum;
                twoDimWrite(QK_t, i, j, N, val);
                }
            }
            //loop over Sequence Length ikj calculate O(i,j) = QK_t(i,k)*V(k,j)
            for(int is = 0; is < N; is += bsize_H){
                for(int js = 0; js < d; js += bsize_W){
                    for(int ks = 0; ks < N; ks += bsize_W){
                        for(int i = is; i < std::min(is+bsize_H, N); i++){
                            for(int k = ks; k < std::min(ks+bsize_W ,N); k++){
                                for(int j = js; j < std::min(js+bsize_W, d); j++){
                                    float val = fourDimRead(O, b, h, i, j, H, N, d);
                                    val += twoDimRead(QK_t, i, k, N) * fourDimRead(V, b, h, k, j, H, N, d);
                                    fourDimWrite(O, b, h, i, j, H, N, d, val);
                                }
                                
                            }
                        }
                    }
                }
            }
    	}
    }
    // DO NOT EDIT THIS RETURN STATEMENT //
    // It formats your C++ Vector O back into a Tensor of Shape (B, H, N, d) and returns it //
    return torch::from_blob(O.data(), {B, H, N, d}, torch::TensorOptions().dtype(torch::kFloat32)).clone();
}

仍然取N = 1000，测试结果如下：

REFERENCE - BLOCKED MATMUL + UNFUSED SOFTMAX statistics
cpu time:  148.424ms
mem usage:  4512000 bytes
    
STUDENT - BLOCKED MATMUL + UNFUSED SOFTMAX statistics
cpu time:  99.278ms
mem usage:  4512000 bytes

比起之前，原reference版本有性能提升，但是自己的版本没有，猜测和计算方式有关。不过无论如何都比reference更快，就不做修改了。

Part 3: Fused Attention

在计算过程中要先计算一个矩阵乘法，再进行softmax，最后再进行一个矩阵乘法，这对缓存性能表现很不利。

解决这个问题的办法是融合计算，这样就只需要一个Nx1的临时向量，而不是NxN的矩阵。在计算Q*K^T的一行后，可以直接进行softmax，然后乘上V，只需要重复这个过程完成计算。在以上基础上还可以进行并行，使用OpenMP进行多线程计算。

融合优化后的代码：

torch::Tensor myFusedAttention(torch::Tensor QTensor, torch::Tensor KTensor, torch::Tensor VTensor, torch::Tensor temp,
                int B, int H, int N, int d){

    // Q, K, V are passed in with Shape: (B, H, N, d)

    //Make O Tensor with Shape (B, H, N, d)
    //and O Row Tensor with Shape (N)
    at::Tensor OTensor = at::zeros({B, H, N, d}, at::kFloat);
    at::Tensor ORowTensor = at::zeros({N}, at::kFloat);

    //Format Y, Q, K, and V tensors into 4D vectors
    std::vector<float> O = formatTensor(OTensor);
    std::vector<float> Q = formatTensor(QTensor);
    std::vector<float> K = formatTensor(KTensor);
    std::vector<float> V = formatTensor(VTensor);
    
    //Format ORow Tensor into a 1D vector
    // You can simply access this as ORow[i]
    std::vector<float> ORow = formatTensor(ORowTensor);

    // -------- YOUR CODE HERE  -------- //
    // We give you a template of the first three loops for your convenience
    //loop over batch
    #pragma omp parallel for collapse(3)
    for (int b = 0; b < B; b++){
        //loop over heads
        for (int h = 0; h < H; h++){
            for (int i = 0; i < N ; i++){
		        // YRow is moved inside so each OpenMP thread gets a local copy.
                at::Tensor ORowTensor = temp.index({torch::indexing::Slice(omp_get_thread_num(), torch::indexing::None)});      
                std::vector<float> ORow = formatTensor(ORowTensor);
                std::fill(ORow.begin(), ORow.end(), 0);
                float sum = 0.0;
                for(int j = 0; j < N; j++){
                    float val = 0;
                    for(int k = 0; k < d; k++){
                        val += fourDimRead(Q, b, h, i, k, H, N, d) * fourDimRead(K, b, h, j, k, H, N, d);
                    }
                    float tmp = exp(val);
                    ORow[j] = tmp;
                    sum += tmp;
                }
                //softmax
                for(int j = 0; j < N; j++) ORow[j] =  ORow[j]/sum;
                //calculate O
                for(int k = 0; k < N; k++){
					for(int j = 0; j < d; j++){
                        float val = fourDimRead(O, b, h, i, j, H, N, d);
						val += ORow[k] * fourDimRead(V, b, h, k, j, H, N, d);
                        fourDimWrite(O, b, h, i, j, H, N, d, val);
                    }
				}
            }
	    }
    }
    // DO NOT EDIT THIS RETURN STATEMENT //
    // It formats your C++ Vector O back into a Tensor of Shape (B, H, N, d) and returns it //
    return torch::from_blob(O.data(), {B, H, N, d}, torch::TensorOptions().dtype(torch::kFloat32)).clone();
}

测试结果：

REFERENCE - FUSED ATTENTION statistics
cpu time:  36.62ms
mem usage:  544000 bytes

STUDENT - FUSED ATTENTION statistics
cpu time:  20.223ms
mem usage:  544000 bytess

Part 4 : Putting it all Together - Flash Attention

由于几个原因，注意力公式的融合非常困难。请注意，该公式包括一个矩阵乘法，然后是softmax的逐行计算，最后是另一个矩阵乘法。将这三个运算融合成块之所以困难，是因为softmax必须对整行进行运算。因此，如果想绕过这种依赖关系，就必须跳出框架。这就是Flash Attention的作用所在。

最后这一part主要是按照给出的伪代码完成Flash Attention。因为时间有限且这个part意义不大，就不做了。