अगर आपने कभी किसी LLM को लंबा जवाब जनरेट करते हुए देखा है और सोचा है कि शुरुआती प्रॉम्प्ट फ्लैश के बाद यह इतना धीमा क्यों हो जाता है, तो आप असल में एक हार्डवेयर बॉटलनेक (hardware bottleneck) को लाइव देख रहे हैं। अधिकांश डेवलपर्स अपने पायथन (Python) कोड, फ्रेमवर्क, या मॉडल के विशाल आकार को दोष देते हैं। वे फंक्शन्स को प्रोफाइल करते हैं, ऑप्टिमाइज़र बदलते हैं, और प्रीप्रोसेसिंग से मिलीसेकंड्स बचाने की कोशिश करते हैं। इनमें से कुछ भी असली समस्या को हल नहीं करता। स्पीड लिमिट आपके सॉफ़्टवेयर में नहीं है। यह सिलिकॉन में है।
हर लार्ज लैंग्वेज मॉडल इन्फरेंस जॉब आपके सर्वर में लगे GPU की दो भौतिक विशेषताओं पर निर्भर करती है: यह कितनी तेज़ी से नंबरों को क्रंच (crunch) कर सकता है, और यह उन नंबरों को क्रंच करने के लिए कितनी तेज़ी से सही जगह पर पहुँचा सकता है।
गणित सस्ता है। डेटा को मूव करना महंगा है
GPU मार्केटिंग में कंप्यूट (compute) के बारे में बात करना बहुत पसंद किया जाता है। प्रति सेकंड ट्रिलियन फ्लोटिंग-पॉइंट ऑपरेशन्स। ये संख्याएँ चौंकाने वाली हैं। लेकिन कंप्यूट केवल आधी कहानी है। दूसरी आधी कहानी है मेमोरी बैंडविड्थ (memory bandwidth), वह दर जिस पर डेटा हाई-बैंडविड्थ मेमोरी से उन कंप्यूट कोर्स में जाता है जहाँ वास्तविक अंकगणित (arithmetic) होता है।
एक LLM इन दोनों कड़ियों में से जो भी कमजोर हो, उससे तेज़ नहीं चल सकता। एक कमर्शियल किचन की कल्पना करें जिसमें बीस मास्टर शेफ हैं। ओवन गर्म हैं, चाकू तेज़ हैं, और हर रसोइया तैयार है। लेकिन सामान की डिलीवरी साइकिल से हो रही है, एक बार में एक टोकरी। किचन रुक जाता है। ज़्यादा शेफ जोड़ने से यह ठीक नहीं होगा। तेज़ ओवन खरीदने से भी यह ठीक नहीं होगा। बॉटलनेक सड़क है।
आधुनिक डेटासेंटर GPU में, अरिथमेटिक यूनिट्स इतनी शक्तिशाली होती हैं कि वे अक्सर अपनी गणनाएँ पूरी कर लेती हैं और फिर खाली बैठी रहती हैं, जबकि वे मेमोरी से वेट्स (weights) और एक्टिवेशन्स (activations) के स्ट्रीम होने का इंतज़ार करती हैं। यह असंतुलन आपके कोड में कोई बग नहीं है। यह चिप्स के निर्माण की भौतिक वास्तविकता है। मेमोरी बैंडविड्थ कच्चे कंप्यूट (raw compute) की गति के साथ तालमेल नहीं बिठा पाई है, और LLMs इस असंतुलन के प्रति विशेष रूप से क्रूर हैं क्योंकि उनके फॉरवर्ड पास (forward passes) में हर सिंगल आउटपुट टोकन के लिए हर सिंगल पैरामीटर को छूना पड़ता है।
प्रॉम्प्ट तेज़ क्यों लगते हैं और जनरेशन धीमा क्यों महसूस होता है
LLM इन्फरेंस दो अलग-अलग चरणों में विभाजित होता है, और वे हार्डवेयर पर पूरी तरह से अलग तरह से दबाव डालते हैं।
Prefill तब होता है जब आपका प्रॉम्प्ट पहली बार मॉडल तक पहुँचता है। सभी टोकन एक साथ आते हैं। GPU बड़े मैट्रिक्स-मैट्रिक्स मल्टीप्लिकेशन का उपयोग करके उन्हें समानांतर (parallel) रूप से प्रोसेस कर सकता है। हज़ारों अरिथमेटिक यूनिट्स एक साथ काम करती हैं, और वर्कलोड घना बना रहता है। यह चरण कंप्यूट-बाउंड (compute-bound) है। शुरुआत में जो आपको गति का अचानक उछाल दिखता है? वह GPU द्वारा वही काम करना है जिसके लिए उसे बनाया गया है।
Decode वह चरण है जहाँ चीज़ें कठिन हो जाती हैं। जब मॉडल अगला टोकन जनरेट करता है, तो वह इसे एक बार में एक टोकन करके करता है। यह चरण मैट्रिक्स-वेक्टर ऑपरेशन्स पर निर्भर करता है, जो GPU की समानांतर क्षमता के केवल एक बहुत छोटे हिस्से का उपयोग करते हैं। इससे भी बुरा यह है कि हर नया टोकन GPU को मेमोरी से पूरे मॉडल वेट्स को फिर से लोड करने के लिए मजबूर करता है। अरिथमेटिक यूनिट्स काम करना चाहती हैं। इसके बजाय, वे इंतज़ार करती हैं। डिकोड मेमोरी-बाउंड (memory-bound) है। GPU प्रभावी रूप से एक महंगे ट्रैफिक कंट्रोलर की तरह काम कर रहा है, जो मैथ इंजन के ठंडा होने तक मेमोरी बस के माध्यम से पैरामीटर्स को इधर-उधर भेजता रहता है। यही कारण है कि सौ शब्दों का जवाब दस सेकंड ले सकता है, भले ही शुरुआती प्रॉम्प्ट विश्लेषण तुरंत महसूस हुआ हो।
KV कैश इसे और भी दिलचस्प बनाता है। डिकोड के दौरान, मॉडल हर पिछले टोकन के लिए की (key) और वैल्यू (value) टेंसर को स्टोर करता है ताकि उसे अटेंशन (attention) को फिर से शुरू से न गिनना पड़े। वह कैश सीक्वेंस की लंबाई के साथ बढ़ता जाता है। यह भी मेमोरी में रहता है। इसलिए अब GPU न केवल वेट्स को फिर से लोड कर रहा है; यह हर सिंगल फॉरवर्ड पास पर लगातार बढ़ते कैश को पढ़ और लिख रहा है। कंप्यूट कोर्स मुश्किल से पसीना बहा रहे हैं, जबकि मेमोरी बस उन दोनों के लिए पसीना बहा रही है।
मेमोरी वॉल (Memory Wall) से लड़ना
इंजीनियरों ने इस बात को कम करने के लिए तकनीकों का एक छोटा शस्त्रागार विकसित किया है कि कितना डेटा मूव करना होगा, या कम से कम डेटा मूव करने की लागत को साझा करने के लिए।
Batching सबसे सीधा तरीका है। यदि एक उपयोगकर्ता का अनुरोध मेमोरी से पूरे वेट लोड को मजबूर करता है, तो एक साथ आठ या सोलह अनुरोधों को प्रोसेस करने से GPU उस लोड को उन सभी में बाँट (amortize) सकता है। वेट्स को एक बार पढ़ा जाता है और बैच के हर सीक्वेंस के लिए पुन: उपयोग किया जाता है। प्रोडक्शन में, परिष्कृत शेड्यूलिंग सिस्टम अनुरोधों को गतिशील रूप से समूहित करते हैं, जिसे कभी-कभी continuous या in-flight batching कहा जाता है, ताकि GPU को शायद ही कभी रुकना पड़े। यह एक ही रूट पर चलने वाली एक बस और सोलह अलग-अलग कारों के बीच के अंतर जैसा है।
Quantization सीधे तौर पर बैंडविड्थ की समस्या पर प्रहार करता है। मॉडल वेट्स (weights) आमतौर पर सोलह-बिट फ्लोटिंग-पॉइंट फॉर्मेट में स्टोर किए जाते हैं। उन्हें आठ-बिट या यहाँ तक कि चार-बिट पूर्णांकों (integers) में कंप्रेस करके, आप बस (bus) के माध्यम से यात्रा करने वाले डेटा की मात्रा को वास्तव में आधा या उससे अधिक कम कर देते हैं। मॉडल को सुसंगत आउटपुट देने के लिए अभी भी पर्याप्त सटीकता (precision) की आवश्यकता होती है, लेकिन आधुनिक पोस्ट-ट्रेनिंग क्वांटाइजेशन विधियाँ गुणवत्ता को खराब किए बिना मॉडल के मेमोरी फुटप्रिंट को नाटकीय रूप से कम कर सकती हैं। डेटा का प्रवाह (data in flight) कम होने का मतलब है मेमोरी कंट्रोलर पर प्रतीक्षा में बिताया गया कम समय।
FlashAttention अटेंशन मैकेनिज्म को इस तरह से पुनर्गठित (restructure) करता है कि मध्यवर्ती परिणाम (intermediate results) GPU की तेज़ ऑन-चिप मेमोरी के भीतर ही रहें। स्टैंडर्ड अटेंशन को बड़े अटेंशन मैट्रिसेस को धीमी बाहरी मेमोरी में लिखना पड़ता था और फिर उन्हें वापस पढ़ना पड़ता था। FlashAttention गणना को छोटे टाइल्स (tiles) में तोड़ देता है जो SRAM में फिट हो जाते हैं, ऑन-चिप पर ही सॉफ्टमैक्स और स्केलिंग स्टेप्स को पूरा करता है, और केवल अंतिम आउटपुट को ही हाई-बैंडविड्थ मेमोरी में वापस लिखता है। यह मुख्य मेमोरी के बहुत कम राउंड ट्रिप्स के बदले थोड़ा अतिरिक्त कंप्यूट (compute) का उपयोग करता है, जो लगभग हमेशा एक फायदेमंद सौदा होता है।
PagedAttention मेमोरी की बर्बादी के एक अलग प्रकार को हल करता है। डिकोड के दौरान, KV कैश अप्रत्याशित रूप से बढ़ता है। पारंपरिक सिस्टम प्रत्येक सीक्वेंस के लिए मेमोरी के निश्चित, निरंतर (contiguous) चंक्स आवंटित करते हैं, जिससे कुछ सीक्वेंस जल्दी समाप्त होने और अन्य के विस्तार होने पर बड़े अंतराल (holes) रह जाते हैं। PagedAttention ऑपरेटिंग सिस्टम से वर्चुअल मेमोरी की अवधारणा उधार लेता है। यह KV कैश प्रविष्टियों को निश्चित आकार के ब्लॉक्स में स्टोर करता है जिन्हें नॉन-कंटीगुअसली (non-contiguously) आवंटित किया जा सकता है और एक इंडिरेक्शन टेबल (indirection table) के माध्यम से मैप किया जा सकता है। यह मेमोरी को आरक्षित लेकिन आधे खाली बफ़र्स के भीतर खाली बैठने से रोकता है और बड़े बैच साइज की अनुमति देता है, जो बदले में फ्रैगमेंटेशन ओवरहेड के बजाय मेमोरी बस को उपयोगी काम में व्यस्त रखकर समग्र थ्रूपुट (throughput) में सुधार करता है।
सवाल बदलें
जब लेटेंसी (latency) बढ़ती है, तो बहुत सी टीमें पूछती हैं कि क्या उन्हें छोटे मॉडल पर स्विच करना चाहिए या अपने इन्फरेंस सर्वर को फिर से लिखना चाहिए। वे सवाल महत्वपूर्ण हैं, लेकिन वे माध्यमिक हैं। पहला सवाल हार्डवेयर के बारे में होना चाहिए। क्या आपका GPU वास्तव में कंप्यूटिंग में व्यस्त है, या यह डेटा के लिए तरस रहा है?
अपने यूटिलाइजेशन मेट्रिक्स (utilization metrics) को देखें। GPU कंप्यूट ऑक्यूपेंसी (occupancy) के साथ-साथ मेमोरी बैंडविड्थ सैचुरेशन का प्रोफाइलिंग करें। यदि आप डिकोड के दौरान उच्च मेमोरी कंटेंशन (contention) और कम अरिथमेटिक इंटेंसिटी (arithmetic intensity) देखते हैं, तो यह मॉडल आर्किटेक्चर की समस्या नहीं है। यह फिजिक्स की समस्या है। समाधान साफ-सुथरे Python कोड से नहीं आएगा। यह अधिक आक्रामक रूप से बैचिंग करने, पाइप के माध्यम से तेज़ी से निकलने के लिए अपने वेट्स को क्वांटाइज़ करने, ऑन-चिप रहने के लिए अटेंशन को पुनर्गठित करने, और KV कैश को प्रबंधित करने से आएगा ताकि आप जगह खत्म हुए बिना बड़े बैच फिट कर सकें।
एक बार जब आप इन्फरेंस को इस नज़रिए से देखने लगते हैं, तो ऑप्टिमाइज़ेशन यांत्रिक (mechanical) हो जाता है। आप मॉडल इंटेलिजेंस के चीजों को धीमा करने के मिथकों के पीछे भागना बंद कर देते हैं और हार्डवेयर वास्तव में जो दे सकता है, उस पर आधारित इंजीनियरिंग निर्णय लेना शुरू कर देते हैं। यही वह बदलाव है जो स्केल करने वाले प्रोडक्शन सिस्टम को केवल काम करने वाले सिस्टम से अलग करता है।