आखिर कर्नल (kernel) क्या है?

GPU पर दो संख्याएँ

आप पहले से ही जानते हैं कि Python में c = a + b कैसे करते हैं। आपने इसे हज़ारों बार किया है। यहाँ हम इस बारे में PyTorch tensors के संदर्भ में बात करेंगे। Tensor सिर्फ संख्याओं की एक सरणी (array) है। इसे GPU पर रखने का मतलब है कि वह सरणी सामान्य CPU मेमोरी के बजाय GPU की मेमोरी में रहती है। जब a और b GPU पर दो tensors होते हैं, तो वह एक पंक्ति इतनी तेज़ी से पूरी होती है कि आप इसके बारे में कभी नहीं सोचते।

अब इसे छोटा करें। मान लें a और b दो एकल floats हैं, दोनों GPU पर हैं। वही पंक्ति। c = a + b। चिप पर वास्तव में क्या चलता है?

उत्तर है kernel। इस दुनिया में, kernel एक छोटा प्रोग्राम है जिसे GPU कुछ डेटा पर चलाता है। वह OS kernel नहीं जिससे आपका लैपटॉप बूट होता है। वह math kernel नहीं जो रैखिक बीजगणित की पाठ्यपुस्तक में होता है। इस शब्द का बहुत अधिक पुन: उपयोग होता है और यह आपकी गलती नहीं है। GPU की दुनिया में, kernel का मतलब बस है: एक छोटा फ़ंक्शन जिसे GPU को अभी, समानांतर में, आपके द्वारा दिए गए डेटा पर चलाने के लिए कहा जाता है।

इसके अंत तक, आप PyTorch कोड के एक टुकड़े को देखकर गिन सकेंगे कि GPU कितने kernels चलाएगा। यह एक छोटी सी तरकीब लगती है, और यह है, लेकिन यह "GPU को ब्लैक बॉक्स" समझने की भावना को दूर करने की पहली सीढ़ी भी है, वह भावना जहाँ आपका मॉडल धीमा है और आप नहीं जानते क्यों। GPU प्रदर्शन के बारे में आप जो भी प्रश्न पूछ सकते हैं, वे अंततः "कौन से kernels चले, और वे क्या कर रहे थे?" पर आकर टिकते हैं। तो हम यहीं से शुरू करते हैं।

आपका पहला kernel

चलिए a और b को थोड़ा बड़ा बनाते हैं। इस बार लंबाई-8 के tensors। अब भी एक PyTorch पंक्ति: c = a + b।

जब आप इसे चलाते हैं, तो आपका CPU (वह मशीन जो वास्तविक Python कर रही है) GPU से कहता है: अरे, इस डेटा पर यह प्रोग्राम चलाओ। उस निर्देश को launch कहा जाता है। जो launch होता है वह एक kernel है: एक प्रोग्राम, चलने के लिए तैयार। Launches स्वयं सस्ते होते हैं, प्रत्येक माइक्रोसेकंड। जो चीज़ launch के आसपास है (GPU पर जाने वाला डेटा, वापस आने वाले परिणाम), वहीं वास्तविक लागत है, और हम आगे बढ़ते हुए इसे गिनेंगे।

kernel के अंदर, वास्तविक काम छोटे श्रमिकों द्वारा किया जाता है जिन्हें threads कहा जाता है। एक GPU के पास हजारों उपलब्ध होते हैं। हमारे लंबाई-8 के जोड़ के लिए, 8 threads काम उठाते हैं: thread 0 तत्व 0 को संभालता है, thread 1 तत्व 1 को, और इसी तरह thread 7 तक। प्रत्येक thread वही छोटा प्रोग्राम चलाता है: a का एक तत्व पढ़ें, b का मिलता-जुलता तत्व पढ़ें, उन्हें जोड़ें, परिणाम c में लिखें।

(व्यवहार में GPU threads को warps नामक निश्चित आकार के समूहों में लॉन्च करता है, NVIDIA कार्ड्स पर हमेशा 32 threads, और जब आपकी सरणी समान रूप से विभाजित नहीं होती तो अतिरिक्त को मास्क कर देता है। अभी के लिए अनदेखा करना सुरक्षित है।)

तो हमारे पास एक PyTorch पंक्ति, एक launch, एक kernel, 8 threads 8 जोड़ कर रहे हैं। अब गिनते हैं कि वास्तव में चिप पर क्या गुज़रा। जोड़ करने के लिए, प्रत्येक thread को a का अपना तत्व और b का अपना तत्व चाहिए। यानी a की 8 रीड और b की 8 रीड। फिर प्रत्येक thread अपना परिणाम c में लिखता है। यानी 8 राइट।

वे रीड और राइट GPU चिप के बगल में बैठी बड़ी मेमोरी में जाते हैं। डेटासेंटर कार्ड्स (A100, H100) पर उस मेमोरी को HBM (हाई-बैंडविड्थ मेमोरी) कहा जाता है। कंज़्यूमर कार्ड्स (RTX 4090) और Colab-शैली के T4s पर, जिन पर लोग वास्तव में इसे आज़माने की संभावना रखते हैं, इसे GDDR कहा जाता है। किसी भी तरह, यह चिप के बगल में बैठी तेज़ मेमोरी है, और हम इसे सिर्फ GPU मेमोरी कहेंगे। यह तेज़ है, लेकिन मुफ़्त नहीं, और इसकी हर यात्रा की कुछ लागत है।

एक kernel = एक launch = डेटा पर एक पास। kernel अपने अंदर जो कुछ भी करता है, उसके किनारों पर रीड और राइट (इनपुट लाने के लिए GPU मेमोरी की बाहरी यात्राएँ, आउटपुट लिखने के लिए वापसी यात्रा) ही वह हिस्सा है जिसकी लागत है। यही पूरा आकार है।

जब tensors बड़े हो जाते हैं तो इसमें कुछ नहीं बदलता। वही PyTorch पंक्ति, वही एक kernel, बस अधिक threads। यदि a और b प्रत्येक में दस लाख तत्व हैं, तो GPU उसी kernel को threads के बड़े झुंड के साथ लॉन्च करता है। गणित स्केल करता है, स्थानांतरित बाइट्स स्केल करती हैं, मानसिक मॉडल नहीं बदलता। एक पंक्ति, एक kernel।

दो ops के बीच क्या होता है

1c = (a + b).relu()

आप Python को इतना अच्छी तरह जानते हैं कि यह दो ऑपरेशन हैं, एक add और फिर एक relu, एक श्रृंखला में। एक इंटरप्रेटर में, यह दो फ़ंक्शन कॉल हैं। GPU पर, eager PyTorch में, यह दो kernel launches हैं: एक add के लिए, एक relu के लिए। अब तक, कोई आश्चर्य नहीं।

वास्तव में दिलचस्प यह है कि दोनों kernels के बीच क्या होता है।

जब add समाप्त होता है, तो उसे अपना परिणाम कहीं रखना होता है। वह कहीं GPU मेमोरी है। Add एक पूर्ण मध्यवर्ती सरणी (इसे tmp कहें) को मेमोरी में लिखता है। फिर एक पल बाद relu लॉन्च होता है, और उसका पहला काम उसी tmp सरणी को मेमोरी से वापस पढ़ना है। वह पूरी चीज़ पढ़ता है, प्रत्येक तत्व पर relu लागू करता है, परिणाम c में लिखता है।

उन दोनों kernels के लिए मेमोरी ट्रैफ़िक गिनें:

• add: a पढ़ता है, b पढ़ता है, tmp लिखता है। तीन सरणी-आकार के स्थानांतरण।
• relu: tmp पढ़ता है, c लिखता है। दो और।

कुल पाँच सरणी-आकार के स्थानांतरण।

इसकी तुलना पिछले अनुभाग के अकेले लंबाई-8 के add से करें, जो तीन था। श्रृंखला में .relu() जोड़ने से आपको सिर्फ relu की गणना की लागत नहीं हुई। इसने आपको GPU मेमोरी के माध्यम से सरणी की एक पूरी राउंड ट्रिप की लागत दी, क्योंकि tmp को लिखना ही पड़ा ताकि अगला kernel इसे वापस पढ़ सके।

कुछ भी कैश नहीं किया गया। tmp किसी रजिस्टर या तेज़ स्थानीय कैश में नहीं रह सका। वह GPU मेमोरी (धीमी और दूर वाली) में गया और तुरंत वापस आ गया। दोनों kernels एक-दूसरे के लिए अजनबी हैं। उन्हें एकमात्र माध्यम के माध्यम से डेटा सौंपना होता है जिससे वे दोनों बात करना जानते हैं: GPU मेमोरी।

PyTorch ऐसा क्यों करता है? क्योंकि eager मोड में, जब आप a + b लिखते हैं, PyTorch उसे अभी चलाता है। उसे नहीं पता कि आप आगे .relu() कॉल करने वाले हैं। प्रत्येक op को उसके Python लाइन के निष्पादित होते ही डिस्पैच कर दिया जाता है। कोई योजना नहीं, कोई आगे की झाँकी नहीं। हर op अपने आप में खड़ा होता है, एक वास्तविक सरणी उत्पन्न करता है, और उसे मेमोरी के माध्यम से जो कुछ भी आगे आता है उसे सौंप देता है।

यही वह पैटर्न है जिसे पकड़कर रखना है। आपके PyTorch कोड में हर मध्यवर्ती मान शारीरिक रूप से GPU मेमोरी में लिखा जाता है और अगले op द्वारा वापस पढ़ा जाता है। हर एक। "kernel count" वास्तव में यही मापता है। हर अतिरिक्त kernel आपके डेटा का GPU मेमोरी के माध्यम से एक और राउंड ट्रिप है।

Fusion: दो ops, एक kernel

एक ऐसे kernel की कल्पना करें जो यह सब एक ही बार में करता है: a का अपना तत्व पढ़ता है, b का अपना तत्व पढ़ता है, उन्हें जोड़ता है, परिणाम पर relu लागू करता है (सब कुछ kernel के अंदर, प्रति-thread छोटी स्क्रैच स्पेस पर जो कभी चिप नहीं छोड़ती), और उसके बाद ही अंतिम मान c में लिखता है। मध्यवर्ती (a + b) अभी भी मौजूद है, लेकिन केवल kernel के अंदर, प्रत्येक thread के निजी स्क्रैच स्पेस में। यह कभी GPU मेमोरी में नहीं लिखा जाता। tmp, एक वास्तविक सरणी के रूप में, बिल्कुल मौजूद नहीं है।

अब स्थानांतरण गिनें। a की रीड: प्रति तत्व 1। b की रीड: प्रति तत्व 1। c की राइट: प्रति तत्व 1। तीन सरणी-आकार के स्थानांतरण। दो-kernel वाले संस्करण के समान गणित, लेकिन दो कम राउंड ट्रिप।

लंबाई 8 पर, यह एक राउंडिंग एरर है। किसी को परवाह नहीं। लंबाई 1 मिलियन या 100 मिलियन पर, वे अतिरिक्त राउंड ट्रिप रनटाइम का एक बड़ा हिस्सा बन जाते हैं, और वॉल क्लॉक इसे दर्शाता है। ऐसा क्यों है कि मेमोरी ट्रैफ़िक इस तरह हावी हो जाता है, यह इस श्रृंखला के लेख 2 का पूरा विषय है, इसलिए मैं 'क्यों' को यहाँ अकेला छोड़ दूँगा। अभी बात यह है: समान गणित, कम यात्राएँ, व्यवहार में तेज़।

वह तरकीब, जो अलग-अलग kernels होने वाले ops को एक single kernel में संयोजित करती है ताकि मध्यवर्ती को कभी GPU मेमोरी पर जाना न पड़े, का एक नाम है। इसे fusion कहा जाता है। बस इतना ही शब्द है। बस इतना ही विचार है।

अब अजीब हिस्सा। उस संयुक्त kernel को हाथ से लिखना add + relu के लिए आसान लगता है। दो ops। बीच में "compute" की एक पंक्ति। लेकिन वास्तविक PyTorch कोड में दर्जनों ops एक श्रृंखला में जुड़े होते हैं, प्रत्येक के अपने आकार, dtype और broadcasting नियम होते हैं। एक fused kernel लिखना जो यह सब सही ढंग से संभाले, एक वास्तविक इंजीनियरिंग कार्य है। आप आमतौर पर इन नियमित elementwise kernels को हाथ से नहीं लिखेंगे।

अच्छी खबर: PyTorch एक उपकरण के साथ आता है जो आपके लिए, ठीक इस प्रकार के मामले के लिए, स्वचालित रूप से पुनर्लेखन करता है। इसे torch.compile कहा जाता है।

आपने लगभग निश्चित रूप से torch.compile को इस तरह उपयोग होते देखा होगा:

1model = torch.compile(model)

एक पंक्ति। इंटरनेट पर किसी ने आपको बताया कि यह चीज़ों को तेज़ बनाता है। यहाँ बताया गया है कि यह वास्तव में क्या करता है, सरल भाषा में: आपके ऑपरेशनों को एक-एक करके चलाने के बजाय जैसा कि eager मोड करता है, torch.compile आपके फ़ंक्शन द्वारा किए जाने वाले tensor ऑपरेशनों को कैप्चर करता है, उन्हें संयोजित करने के अवसरों की तलाश करता है, और अनुकूलित कोड उत्पन्न करता है। बाद में कॉल जो समान मान्यताओं से मेल खाते हैं, उस कोड का पुन: उपयोग कर सकते हैं।

हमने ऊपर कागज पर जो fusion किया था (add और relu, एक kernel साझा करना, tmp कभी मेमोरी को नहीं छूता), ठीक उसी तरह की चीज़ है जो torch.compile आपके कोड के साथ स्वचालित रूप से करेगा, जब तक ops काफी सरल हों। जब लोग कहते हैं कि torch.compile "PyTorch को तेज़ बनाता है," तो उनका यही एक बड़ा हिस्सा होता है।

उन मामलों के लिए जिन्हें torch.compile अपने आप fuse नहीं कर सकता (कस्टम ops जिन्हें यह नहीं पहचानता, असामान्य reductions, अजीब मेमोरी लेआउट), किसी को अभी भी हाथ से kernel लिखना होगा। इसीलिए Triton और CUDA जैसे उपकरण हैं। एक अलग लेख।

इसे स्वयं देखें

ऊपर सब कुछ कागज पर kernels गिन रहा था। अब उन्हें एक वास्तविक GPU पर गिनने का समय है। यदि आपके पास CUDA GPU वाली कोई मशीन है (एक वर्कस्टेशन, एक Colab नोटबुक, एक क्लाउड इंस्टेंस), तो आप इसे कुछ मिनटों में स्वयं चला सकते हैं।

उपकरण है torch.profiler। यह PyTorch में बनाया गया है। यह बस यह रिकॉर्ड करता है कि आपका कोड चलने के दौरान GPU ने वास्तव में क्या किया, और आपको एक तालिका देता है जिसे आप पढ़ सकते हैं।

चरण 1: eager संस्करण।

दो-ऑप वाली पंक्ति को एक फ़ंक्शन में लपेटें ताकि हमारे पास कॉल करने के लिए कुछ हो:

1import torch
2from torch.profiler import profile, ProfilerActivity
3
4def add_relu(a, b):
5    return (a + b).relu()
6
7a = torch.randn(1_000_000, device="cuda")
8b = torch.randn(1_000_000, device="cuda")
9
10with profile(activities=[ProfilerActivity.CUDA]) as prof:
11    add_relu(a, b)
12    torch.cuda.synchronize()
13
14print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))

torch.cuda.synchronize() सिर्फ यह सुनिश्चित करने के लिए है कि टाइमिंग पढ़ने से पहले GPU ने काम पूरा कर लिया है। GPU काम एसिंक्रोनस रूप से चलता है, और सिंक के बिना आप कभी-कभी वास्तविक kernel काम के बजाय launch ओवरहेड को माप सकते हैं।

चरण 2: आउटपुट पढ़ें।

आपका वास्तविक profiler आउटपुट में आपकी अपेक्षा से अधिक पंक्तियाँ होंगी। मेमोरी-आवंटन और PyTorch बुककीपिंग की बहुत सी पंक्तियाँ मिश्रित होंगी। हमारे लिए महत्वपूर्ण पंक्तियाँ CUDA kernels हैं, वे वास्तविक फ़ंक्शन जो GPU ने चलाए। नाम में kernel वाली पंक्तियाँ देखें। वे दो लगभग इस तरह दिखेंगी:

1vectorized_elementwise_kernel<...BinaryFunctor_add...>    1    12us
2vectorized_elementwise_kernel<...threshold_kernel...>     1     9us

सटीक टेम्पलेट नाम PyTorch संस्करणों के बीच बदलते हैं (relu अक्सर threshold के रूप में दिखाई देता है क्योंकि यह अंतर्निहित op है, और add कभी-कभी CUDAFunctor_add के रूप में)। पूरी चीज़ को पार्स करने की कोशिश न करें। बस पंक्तियाँ गिनें। दो पंक्तियाँ। दो kernels। एक add के लिए, एक relu के लिए। बिल्कुल वैसा ही जैसा हमने एक सेक्शन पहले कहा था कि होगा।

चरण 3: compiled संस्करण।

एक पंक्ति बदलें। फ़ंक्शन को torch.compile में लपेटें:

1compiled = torch.compile(add_relu)

अब यहाँ एक बात है जो जानना उपयोगी है इसे चलाने से पहले। torch.compile किए गए फ़ंक्शन की पहली कॉल धीमी होती है, कभी-कभी अजीब तरह से धीमी, क्योंकि तब torch.compile अपना काम करता है: आपके कोड का विश्लेषण करना, यह पता लगाना कि क्या fuse करना है, fused kernel उत्पन्न करना। यदि आप पहली कॉल को प्रोफाइल करते हैं, तो आप compile चरण को माप रहे हैं, kernel को नहीं। तो पैटर्न है: इसे एक बार गर्म करने के लिए कॉल करें, परिणाम को फेंक दें, फिर प्रोफाइल करें।

1compiled(a, b)                # वार्म-अप, फेंक दें
2torch.cuda.synchronize()
3
4with profile(activities=[ProfilerActivity.CUDA]) as prof:
5    compiled(a, b)
6    torch.cuda.synchronize()
7
8print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))

चरण 4: आउटपुट फिर से पढ़ें।

इस बार:

1triton_poi_fused_add_relu_0    1    14us

एक पंक्ति। एक kernel। नाम भी बताता है कि इसने क्या किया: fused add और relu। पहले जैसा ही गणित, दो के बजाय एक launch।

आपने वही किया जिसके बारे में यह लेख बात कर रहा था, एक वाक्य में: आपने PyTorch से दो ops को एक kernel में संयोजित करने के लिए कहा, profiler को देखा, और पुष्टि की कि गिनती दो से एक हो गई। Fusion, वास्तविक जीवन में, आपकी मशीन पर।

यदि आप इसे और अधिक नाटकीय रूप से देखना चाहते हैं, तो इसे कुछ अलग-अलग tensor आकारों पर आज़माएँ। लंबाई 100 पर, दोनों संस्करण इतनी तेज़ी से चलते हैं कि अंतर शोर में खो जाता है। लंबाई 10 मिलियन या 100 मिलियन पर, compiled संस्करण स्पष्ट रूप से आगे निकलने लगता है, क्योंकि हमने जो राउंड ट्रिप काटा है वह उस पैमाने पर काम का एक वास्तविक हिस्सा है।

कर्नेल गिनना अब अमूर्त सलाह नहीं है। आपके पास जाँच करने का एक तरीका है।

चलिए इसे समाप्त करते हैं

यहाँ पूरी बात एक बार में है।

आपका PyTorch कोड, जब GPU पर चलता है, तो kernels के एक अनुक्रम में बदल जाता है। प्रत्येक kernel एक launch है, आपके डेटा पर एक पास, इनपुट लाने और आउटपुट लिखने के लिए GPU मेमोरी के माध्यम से एक राउंड ट्रिप। सरल ops एक kernel बन जाते हैं। ops की श्रृंखलाएँ डिफ़ॉल्ट रूप से per op एक kernel बन जाती हैं, बीच में मध्यवर्ती मेमोरी के माध्यम से राउंड-ट्रिप करते हैं। torch.compile आपके लिए सरल श्रृंखलाओं को fuse कर सकता है ताकि वे मध्यवर्ती कभी मेमोरी को स्पर्श न करें। कम kernels का आमतौर पर मतलब कम मेमोरी ट्रैफ़िक है। और कम मेमोरी ट्रैफ़िक का आमतौर पर मतलब तेज़ है।