PyTorch में प्रोफाइलिंग (भाग 3): प्रोफाइल ही सब कुछ है

PyTorch में प्रोफाइलिंग श्रृंखला का तीसरा पोस्ट अटेंशन मैकेनिज्म को प्रोफाइल करता है। यह एक सामान्य अटेंशन मॉड्यूल (matmul, स्केलिंग, मास्किंग, softmax, matmul) से शुरू होता है और इसका ट्रेस प्रति फॉरवर्ड 5 GPU कर्नेल दिखाता है। out-of-place masked_fill को in-place masked_fill_ से बदलने से एक मेमोरी कॉपी कर्नेल खत्म होता है, जिससे समय और मेमोरी बचती है। इसके बाद, स्केल्ड डॉट प्रोडक्ट अटेंशन (SDPA) प्रस्तुत किया गया है; इसका गणित बैकएंड, हालांकि सुरक्षित (NaN मामलों को संभालना) और अधिक सटीक (FP32) है, सामान्य अटेंशन से 3.7 गुना धीमा है, क्योंकि यह 20 GPU कर्नेल लॉन्च करता है, टेंसर कोर का उपयोग नहीं करता (bf16 टेंसर कोर matmul के बजाय sgemm का उपयोग करता है), हर कॉल में कॉज़ल मास्क को फिर से बनाता है, और _safe_softmax का उपयोग करता है। पोस्ट में नोट किया गया है कि SDPA सामान्यतः स्वचालित रूप से सबसे तेज़ बैकएंड का चयन करता है।