Ada-MK: Adaptive MegaKernel Optimization via Automated DAG-based Search for LLM Inference
Ada-MK एक नवीन अनुकूलन ढांचा (optimization framework) है जो कंपाइल-टाइम DAG-आधारित खोज और मेमोरी स्प्लिटिंग के माध्यम से रनटाइम डायनेमिक शेड्यूलिंग को समाप्त करता है और साझा मेमोरी के उपयोग को कम करता है, जिससे NVIDIA GPUs पर vLLM की तुलना में 50.2% तक थ्रूपुट सुधार प्राप्त करने के लिए वाणिज्यिक ऑनलाइन विज्ञापन प्रणालियों में मेगाकर्नेल (MegaKernels) का पहला औद्योगिक परिनियोजन सक्षम होता है।
मूल पेपर CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) के तहत लाइसेंस किया गया है। नीचे दिए गए पेपर की यह व्याख्या AI से तैयार की गई है। इसे लेखकों ने न तो लिखा है, न इसका समर्थन किया है। तकनीकी सटीकता के लिए मूल पेपर देखें। पूरा डिस्क्लेमर पढ़ें
कल्पना कीजिए कि आप एक विशाल ऑनलाइन स्टोर (जैसे कि एक सर्च इंजन या विज्ञापन प्लेटफॉर्म) के लिए एक हाई-स्पीड डिलीवरी सर्विस चला रहे हैं। हर बार जब कोई ग्राहक कोई सवाल पूछता है, तो आपके सिस्टम को शब्द-दर-शब्द जवाब जेनरेट करना होता है। लार्ज लैंग्वेज मॉडल्स (LLMs) की दुनिया में, इसे "इन्फरेंस" (inference) कहा जाता है।
समस्या यह है कि हर एक शब्द जेनरेट करने के लिए, आपके कंप्यूटर को अपने ग्राफिक्स कार्ड (GPU) को हजारों छोटे निर्देश भेजने पड़ते हैं। यह ऐसा है जैसे एक डिलीवरी ड्राइवर को हर एक पैकेज के लिए पोस्ट ऑफिस रुकना पड़ता है, एक पैकेज उठाना पड़ता है, गोदाम तक जाना पड़ता है, सामान छोड़ना पड़ता है, वापस आना पड़ता है, और सिर्फ एक वाक्य के लिए यह प्रक्रिया हजारों बार दोहरानी पड़ती है। यह "रुकने और चलने" वाला ओवरहेड (stop-and-go overhead) बहुत सारा समय बर्बाद करता है—कुल यात्रा समय का लगभग 15% हिस्सा तो सिर्फ इन छोटे-छोटे स्टॉप्स में निकल जाता है, न कि वास्तव में सामान पहुँचाने में।
बड़ा विचार: "मेगाकर्नेल" (The MegaKernel)
इस समस्या को हल करने के लिए, शोधकर्ताओं ने मेगाकर्नेल की अवधारणा का आविष्कार किया। हर छोटे काम के लिए पोस्ट ऑफिस रुकने के बजाय, कल्पना कीजिए कि एक सुपर-एफिशिएंट डिलीवरी ट्रक है जो शुरुआत में ही वह सब कुछ ले लेता है जिसकी उसे आवश्यकता है, पूरे रास्ते में बिना रुके चलता है, और अंत में सब कुछ एक साथ डिलीवर कर देता है। यह उन सभी छोटे-छोटे स्टॉप्स को एक लंबी, निरंतर यात्रा में जोड़ देता है। इससे "रुकने और चलने" वाली देरी खत्म हो जाती है।
हालांकि, इसमें एक पेंच है। आपकी कंपनी (NVIDIA Ada/L20) द्वारा उपयोग किया जाने वाला विशिष्ट प्रकार का ट्रक (GPU) छोटा है और नए, अधिक शानदार ट्रकों (Hopper/Blackwell) की तुलना में इसके केबिन में स्टोरेज स्पेस कम है।
- पुराना समाधान 1 (हैंड-ट्यून्ड): विशेषज्ञों ने विशेष रूप से इस छोटे केबिन के लिए एक कस्टम ट्रक बनाया। यह तेज़ था, लेकिन अगर आप कार्गो (AI मॉडल) या सड़क (हार्डवेयर) बदलते, तो ट्रक टूट जाता। यह पोर्टेबल नहीं था।
- पुराना समाधान 2 (ऑटो-ट्यून्ड): उन्होंने एक ऐसा ट्रक बनाने की कोशिश की जो चलते-फिरते अपने कार्गो स्पेस को खुद एडजस्ट कर सके। लेकिन ड्राइवर को लगातार मैप देखना पड़ता था और निर्णय लेने पड़ते थे ("क्या शेल्फ भर गई है? क्या मुझे रुकना चाहिए?")। ये लगातार निर्णय लेने की प्रक्रिया ट्रक को धीमा कर देती थी, जो कि अस्वीकार्य है जब आपको मिलीसेकंड में डिलीवरी करनी हो।
नया समाधान: Ada-MK
लेखकों ने Ada-MK बनाया, जो एक नया सिस्टम है जो छोटे "Ada" ट्रकों के लिए इन समस्याओं को हल करता है। उन्होंने इसे कैसे किया, इसके सरल उदाहरण यहाँ दिए गए हैं:
1. "स्मार्ट पैकिंग" रणनीति (Adaptive Shared Memory)
छोटे ट्रक में एक छोटा स्टोरेज कंपार्टमेंट (Shared Memory) होता है। यदि आप इसमें बहुत अधिक सामान भरने की कोशिश करते हैं, तो ट्रक रुक जाता है।
- नवाचार: पूरा सूटकेस छोटे कंपार्टमेंट में डालने के बजाय, उन्होंने सूटकेस को आधा काट दिया (K-dimension splitting)। वे केवल आधे सामान को पैक करते हैं, उसे डिलीवर करते हैं, फिर दूसरे आधे हिस्से को पैक करते हैं।
- परिणाम: इससे पीक स्टोरेज की आवश्यकता 50% कम हो गई। उन्होंने यह भी पता लगाया कि कैसे एक पैकेज छोड़ने के तुरंत बाद खाली जगह का पुन: उपयोग किया जाए, जिससे यह सुनिश्चित हो सके कि ट्रक स्पेस की प्रतीक्षा में कभी भी खाली न बैठा रहे।
2. "पूर्व-नियोजित मार्ग" (Offline DAG Search)
पुराने "ऑटो-ट्यून्ड" ट्रक ड्राइविंग के दौरान निर्णय लेते थे, जिससे ट्रैफिक जाम (branch penalties) होता था।
- नवाचार: टीम ने एक शक्तिशाली कंप्यूटर का उपयोग करके ट्रक के गैरेज से निकलने से पहले ही लाखों संभावित मार्गों का सिमुलेशन किया। उन्होंने हर मोड़ और स्टॉप का एक विस्तृत आरेख (DAG) तैयार किया।
- परिणाम: एक बार सबसे अच्छा मार्ग मिल जाने के बाद, उन्होंने इसे "ठोस" (solidify) कर दिया। ड्राइवर को अब ड्राइविंग के दौरान सोचने या मैप चेक करने की आवश्यकता नहीं है; वे बस पूर्व-नियोजित पथ का पालन करते हैं। यह निर्णय लेने की देरी को हटा देता है, जिससे ड्राइव अविश्वसनीय रूप से सुचारू और तेज़ हो जाती है।
3. "हाइब्रिड फ्लीट" (Heterogeneous Engine)
उन्होंने महसूस किया कि यात्रा के कुछ हिस्सों के लिए (शुरुआत में भारी मात्रा में पैकेज लोड करना, जिसे "प्रीफिल" कहा जाता है), पुराने मानक ट्रक वास्तव में बेहतर थे। लेकिन अंतिम डिलीवरी (शब्द-दर-शब्द जेनरेट करना, जिसे "डिकोड" कहा जाता है) के लिए, उनका नया मेगाकर्नेल ट्रक श्रेष्ठ था।
- नवाचार: उन्होंने एक हाइब्रिड सिस्टम बनाया। उन्होंने भारी कामों के लिए मानक ट्रकों को बनाए रखा, लेकिन फिर निर्बाध रूप से अपने नए मेगाकर्नेल ट्रक को अंतिम डिलीवरी चरण के लिए बदल दिया।
- परिणाम: आपको दोनों दुनियाओं का सर्वश्रेष्ठ मिलता है: शुरुआत के लिए उच्च गति और अंतिम डिलीवरी के लिए अल्ट्रा-लो लेटेंसी, बिना पूरे डिलीवरी नेटवर्क को फिर से बनाए।
वास्तविक दुनिया का प्रभाव
टीम ने इसे बैडू (Baidu) के कमर्शियल ऑनलाइन विज्ञापन सिस्टम पर टेस्ट किया।
- गति: उन परिदृश्यों में जहाँ गति महत्वपूर्ण है (छोटे बैच, छोटे रिस्पॉन्स), उनका सिस्टम मानक उद्योग उपकरणों की तुलना में 23.6% तक तेज़ और vLLM नामक एक लोकप्रिय ओपन-सोर्स टूल की तुलना में 50.2% तेज़ था।
- विश्वसनीयता: यह विभिन्न प्रकार के AI मॉडल और कार्यों के लिए लगातार काम करता रहा।
- प्रथम श्रेणी का: यह पहली बार है जब एक "मेगाकर्नेल" को सफलतापूर्वक एक लाइव कमर्शियल एडवरटाइजिंग सिस्टम में तैनात किया गया है।
सारांश में
यह पेपर बताता है कि विशिष्ट, छोटे कंप्यूटर चिप्स पर AI चैटबॉट्स और विज्ञापन सिस्टम को बहुत तेज़ कैसे बनाया जा सकता है। उन्होंने इसे डेटा को अधिक कुशलता से पैक करके, पूरे डिलीवरी रूट को पहले से प्लान करके ताकि ड्राइवर को कभी सोचना न पड़े, और अपने नए हाई-स्पीड डिलीवरी मेथड को मौजूदा टूल्स के साथ मिलाकर किया। परिणाम स्वरूप, एक ऐसा सिस्टम मिलता है जो जवाबों को काफी तेज़ी से डिलीवर करता है, खासकर जब कई उपयोगकर्ता एक-एक करके सवाल पूछ रहे हों।
अपने क्षेत्र के पेपरों की भीड़ में उलझे हुए हैं?
आपके रिसर्च कीवर्ड से मेल खाने वाले सबसे नए और अलग सोच वाले पेपरों का रोज़ाना Digest पाएँ—तकनीकी सारांश के साथ, आपकी भाषा में।