FlashAttention-2: बेहतर parallel processing और work partitioning से और तेज़ हुआ Attention

• GPT-4(32k), MPT(65k), Calude(100k) जैसे लंबे context वाले language model सामने आए हैं
• Transformer की context length बढ़ाना मुश्किल है क्योंकि runtime और memory requirements चौथी घात तक बढ़ जाती हैं
• पिछले साल जारी किया गया FlashAttention memory usage घटाकर और attention speed बढ़ाकर कई जगहों पर इस्तेमाल किया जा रहा है
• रिलीज़ के समय ही यह 2-4 गुना तेज़ था, लेकिन अभी भी सुधार की गुंजाइश है. optimized matrix multiplication (GEMM) की तुलना में यह अब भी उतना तेज़ नहीं है और सैद्धांतिक अधिकतम FLOPs/s का केवल 25-40% ही हासिल करता है (A100 GPU पर अधिकतम 124 TFLOPs/s)
• FlashAttention-2 पिछले version से 2 गुना तेज़ है और A100 GPU पर अधिकतम 230 TFLOP/s performance देता है
• GPT प्रकार के language model training में यह अधिकतम 225 TFLOPS तक पहुँचा (72% model FLOP utilization)
• algorithm को समायोजित करके non-matmul FLOPs घटाए गए
• बेहतर parallelization और प्रत्येक thread block में work partitioning के तरीके में बदलाव
• Head Dimensions की संख्या 128 से बढ़ाकर 256 की गई