PTXBench: Benchmark and Adapt LLMs for GPU Kernel Optimization with Architecture-specific PTX

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

शोधकर्ताओं की एक टीम ने PTXBench नामक एक नया परीक्षण क्षेत्र बनाकर इसका उत्तर देने का प्रयास किया। वे यह देखना चाहते थे कि क्या ये AI मॉडल न केवल ऐसा कोड लिख सकते हैं जो काम करता है, बल्कि ऐसा कोड भी जो नवीनतम GPU की विशिष्ट, उन्नत विशेषताओं का सक्रिय रूप से उपयोग करता है। शोधकर्ताओं ने दो सबसे शक्तिशाली ग्राफिक्स कार्डों, H100 और B200 पर ध्यान केंद्रित किया, और उन्होंने मॉडल्स का परीक्षण आर्टिफिशियल इंटेलिजेंस में उपयोग किए जाने वाले जटिल गणितीय कार्यों, जैसे कि मैट्रिक्स मल्टीप्लिकेशन और अटेंशन मैकेनिज्म से जुड़े कार्यों पर किया। लक्ष्य तीन चीजों को मापना था: क्या कोड बिना किसी त्रुटि के चला, क्या इसने वास्तव में उन विशिष्ट हार्डवेयर निर्देशों को ट्रिगर किया जिन्हें शोधकर्ताओं ने मांगा था, और क्या यह मौजूदा सर्वश्रेष्ठ सॉफ्टवेयर लाइब्रेरी से तेज़ चला।

परिणामों ने उस अंतर को उजागर किया जो मॉडल्स की क्षमता और शीर्ष-स्तरीय प्रदर्शन की आवश्यकता के बीच था। जबकि AI मॉडल अक्सर सरल कार्यों के लिए सही ढंग से काम करने वाला कोड लिखने में सफल रहे, उन्हें अधिक जटिल कार्यों के साथ काफी संघर्ष करना पड़ा, विशेष रूप से अटेंशन मैकेनिज्म के 'बैकवर्ड पास' से जुड़े कार्यों के साथ, जो AI सिस्टम को प्रशिक्षित करने के लिए महत्वपूर्ण हैं। भले ही मॉडल्स ने अनुरोधित हार्डवेयर निर्देशों को निष्पादित करने वाला कोड लिखने में सफलता प्राप्त की, लेकिन वह कोड शायद ही कभी अग्रणी पेशेवर लाइब्रेरी की गति से मेल खाता था। कई मामलों में, AI-जनरेटेड कोड धीमा था, जिससे पता चलता है कि केवल निर्देशों को चलाना ही वास्तविक अनुकूलन (ऑप्टिमाइजेशन) प्राप्त करने के लिए पर्याप्त नहीं था। अध्ययन में पाया गया कि मॉडल्स अक्सर अद्वितीय, आर्किटेक्चर-विशिष्ट विशेषताओं का लाभ उठाने के बजाय जेनेरिक कोडिंग पैटर्न पर वापस चले जाते थे, भले ही वे विशेषताएँ उन्हें स्पष्ट रूप से बताई गई थीं।

इन मॉडल्स को सुधारने के तरीके को समझने के लिए, शोधकर्ताओं ने एक नियंत्रित प्रयोग किया जहाँ उन्होंने 'सुपरवाइज्ड फाइन-ट्यूनिंग' नामक विधि का उपयोग करके एक विशिष्ट AI मॉडल को प्रशिक्षित किया। उन्होंने मॉडल को उसकी अपनी गलतियों के उदाहरण, उनके सुधारे गए संस्करणों और इस बात के स्पष्टीकरण के साथ दिखाया कि सुधार क्यों काम कर रहे थे। इस दृष्टिकोण ने, जिसे उन्होंने Fixit कहा, मॉडल को कई कार्यों में सुधार करने में मदद की, विशेष रूप से कार्यात्मक रूप से सही कोड उत्पन्न करने में। हालाँकि, सुधार समान नहीं था। मॉडल अभी भी अलग-अलग समस्या आकारों या उन कार्यों के विविध रूपों पर अपने नए कौशल को सामान्यीकृत (जनरलाइज) करने में संघर्ष करता रहा जो प्रशिक्षण के दौरान उसके सामने नहीं आए थे। शोधकर्ताओं ने पाया कि सुधारों की गुणवत्ता और प्रशिक्षण डेटा का संतुलन उतना ही महत्वपूर्ण था जितना कि डेटा की कुल मात्रा। एक विविध सेट की समस्याओं पर प्रशिक्षित मॉडल, कम विविध लेकिन बड़े सेट पर प्रशिक्षित मॉडल की तुलना में बेहतर प्रदर्शन करता था, जो दर्शाता है कि सीखने का प्रकार मात्रा से अधिक महत्वपूर्ण था।

अध्ययन ने AI को सही संदर्भ प्रदान करने के महत्व पर भी प्रकाश डाला। जब शोधकर्ताओं ने मॉडल्स को हार्डवेयर की विशिष्ट क्षमताओं के बारे में विस्तृत जानकारी दी, तो मॉडल्स द्वारा अनुरोधित निर्देशों का उपयोग करने की संभावना बहुत बढ़ गई। इस विशिष्ट ज्ञान के बिना, मॉडल शायद ही कभी अनुरोधित PTX का उपयोग करता है। यह सुझाव देता है कि जबकि AI इस निम्न स्तर पर प्रोग्रामिंग करना सीख सकता है, इसके लिए सटीक मार्गदर्शन और उच्च-गुणवत्ता वाले उदाहरणों की आवश्यकता होती है। शोधकर्ताओं ने निष्कर्ष निकाला कि हालांकि वर्तमान AI मॉडल्स आशाजनक संकेत दिखाते हैं, लेकिन वे अभी तक नवीनतम हार्डवेयर के लिए कोड को अनुकूलित करने हेतु विशेषज्ञ मानव इंजीनियरों को बदलने के लिए तैयार नहीं हैं। आगे का रास्ता बेहतर प्रशिक्षण डेटा, अधिक संतुलित सीखने के अनुभवों और इन मॉडल्स को विकसित होते कंप्यूटर आर्किटेक्चर की विशिष्ट बाधाओं के बारे में तर्क करने के लिए सिखाने की गहरी समझ से जुड़ा है।