Methods and devices for accelerating a transformer with a sparse attention pattern
A method and an apparatus for accelerating a transformer with a sparse attention pattern are provided. The method includes that a heterogeneous device including one or more GPUs loads a first matrix, a second matrix, and a transformed sparsity mask into a first sampled dense-dense matrix multiplication (SDDMM) kernel in a sparse attention module in the transformer and generates a first output based on the first matrix, the second matrix, and the transformed sparsity mask by the first SDDMM kernel, generates a second output by a softmax kernel in the sparse attention module based on the first output, loads the second output, a third matrix, and the transformed sparsity mask into a matrix multiplication kernel in the sparse attention module, and generates an output of the sparse attention module.
1 . A method for accelerating a transformer with a sparse attention pattern in heterogeneous devices, comprising:
loading, by a heterogeneous device comprising one or more graphic processing units, a first matrix, a second matrix, and a transformed sparsity mask into a first sampled dense-dense matrix multiplication (SDDMM) kernel in a sparse attention module in the transformer and generating a first output based on the first matrix, the second matrix, and the transformed sparsity mask by the first SDDMM kernel;
generating, by the heterogeneous device, a second output by a softmax kernel in the sparse attention module based on the first output; and
loading, by the heterogeneous device, the second output, a third matrix, and the transformed sparsity mask into a matrix multiplication kernel in the sparse attention module and generating an output of the sparse attention module, wherein the matrix multiplication kernel is a second SDDMM kernel.
2 . The method of claim 1 , wherein the transformed sparsity mask and the first output are stored in a compressed format.
3 . The method of claim 2 , further comprising:
transforming, by the heterogeneous device, a sparsity mask in a regular dense matrix format to the transformed sparsity mask in the compressed format, wherein the sparsity mask indicates the sparsity attention pattern.
4 . The method of claim 1 , wherein the first matrix, the second matrix, and the third matrix are in a regular dense matrix format.
5 . The method of claim 2 , wherein loading the second output, the third matrix, and the transformed sparsity mask into the matrix multiplication kernel in the sparse attention module comprises:
loading the second output, the third matrix, and the transformed sparsity mask into the second SDDMM kernel in the sparse attention module to generate the output of the sparse attention module.
6 . The method of claim 1 , further comprising:
scaling and applying, by the softmax kernel, a softmax function over the first output to generate the second output, wherein the first output is a sparse matrix.
7 . An apparatus for accelerating a transformer with a sparse attention pattern in heterogeneous devices, comprising:
one or more processors; and
a memory configured to store instructions executable by the one or more processors,
wherein the one or more processors, upon execution of the instructions, are configured to:
load a first matrix, a second matrix, and a transformed sparsity mask into a first sampled dense-dense matrix multiplication (SDDMM) kernel in a sparse attention module in the transformer and generate a first output based on the first matrix, the second matrix, and the transformed sparsity mask by the first SDDMM kernel;
generate a second output by a softmax kernel in the sparse attention module based on the first output; and
load the second output, a third matrix, and the transformed sparsity mask into a matrix multiplication kernel in the sparse attention module and generate an output of the sparse attention module, wherein the matrix multiplication kernel is a second SDDMM kernel.
8 . The apparatus of claim 7 , wherein the transformed sparsity mask and the first output are stored in a compressed format.
9 . The apparatus of claim 8 , wherein the one or more processors are further configured to:
transform a sparsity mask in a regular dense matrix format to the transformed sparsity mask in the compressed format, wherein the sparsity mask indicates the sparsity attention pattern.
10 . The apparatus of claim 7 , wherein the one or more processors are further configured to:
wherein the first matrix, the second matrix, and the third matrix are in a regular dense matrix format.
11 . The apparatus of claim 8 , wherein the one or more processors are further configured to:
load the second output, the third matrix, and the transformed sparsity mask into the second SDDMM kernel in the sparse attention module to generate the output of the sparse attention module.
12 . The apparatus of claim 7 , wherein the one or more processors are further configured to:
scale and apply, by the softmax kernel, a softmax function over the first output to generate the second output, wherein the first output is a sparse matrix.
13 . A non-transitory computer readable storage medium, comprising instructions stored therein, wherein, upon execution of the instructions by one or more processors, the instructions cause the one or more processors to perform acts comprising:
loading a first matrix, a second matrix, and a transformed sparsity mask into a first sampled dense-dense matrix multiplication (SDDMM) kernel in a sparse attention module in a transformer in heterogeneous devices and generating a first output based on the first matrix, the second matrix, and the transformed sparsity mask by the SDDMM kernel;
generating a second output by a softmax kernel in the sparse attention module based on the first output; and
loading the second output, a third matrix, and the transformed sparsity mask into a matrix multiplication kernel in the sparse attention module and generating an output of the sparse attention module, wherein the matrix multiplication kernel is a second SDDMM kernel.
14 . The non-transitory computer readable storage medium of claim 13 , wherein the transformed sparsity mask and the first output are stored in a compressed format.
15 . The non-transitory computer readable storage medium of claim 14 , wherein the one or more processors are caused to perform acts further comprising:
transforming, by the heterogeneous device, a sparsity mask in a regular dense matrix format to the transformed sparsity mask in the compressed format, wherein the sparsity mask indicates the sparsity attention pattern.
16 . The non-transitory computer readable storage medium of claim 13 , wherein the first matrix, the second matrix, and the third matrix are in a regular dense matrix format.
17 . The non-transitory computer readable storage medium of claim 14 , wherein the one or more processors are caused to perform acts further comprising:
loading the second output, the third matrix, and the transformed sparsity mask into the second SDDMM kernel in the sparse attention module to generate the output of the sparse attention module.
18 . The non-transitory computer readable storage medium of claim 13 , wherein the one or more processors are caused to perform acts further comprising:
scaling and applying, by the softmax kernel, a softmax function over the first output to generate the second output, wherein the first output is a sparse matrix.