𝗞𝗩 𝗖𝗮𝗰𝗵𝗲 𝗮𝗻𝗱 𝗣𝗮𝗴𝗲𝗱𝗔𝘁𝘁𝗲𝗻𝘁𝗶𝗼𝗻: 𝗪𝗵𝘆 𝗬𝗼𝘂𝗿 𝗟𝗟𝗠 𝗦𝗲𝗿𝘃𝗲𝗿 𝗦𝗹𝗼𝘄𝘀 𝗗𝗼𝘄𝗻

आपका LLM सर्वर धीमा चल रहा है।

आपने चार A100 GPUs पर एक 70B मॉडल तैनात किया है। सुबह 8 बजे तक सब कुछ ठीक दिखता है। दोपहर के भोजन तक, लेटेंसी (latency) दोगुनी हो जाती है। आप अपनी मेमोरी चेक करते हैं। इसका अधिकांश हिस्सा "tensor buffers" द्वारा घेरा गया है। ये वास्तव में पुरानी बातचीत के कैश्ड स्टेट्स (cached states) हैं।

यह KV cache की समस्या है। प्रोडक्शन LLM सर्विंग में यह सबसे बड़ी बाधा (bottleneck) है।

KV cache क्या है?

हर transformer मॉडल एक-एक करके टोकन जनरेट करता है। एक नया टोकन बनाने के लिए, मॉडल को पिछले सभी टोकन के Key और Value tensors की आवश्यकता होती है। इन्हें हर बार फिर से कैलकुलेट करना बहुत धीमा होता है। इसके बजाय, इंजन उन्हें स्टोर कर लेता है। यही स्टोरेज KV cache कहलाता है।

मेमोरी की समस्या:

Llama 3.1 70B मॉडल के लिए, एक सिंगल 4096-token सीक्वेंस को लगभग 1.3 GB मेमोरी की आवश्यकता होती है।

यदि एक साथ 256 उपयोगकर्ता हैं, तो आपको 336 GB मेमोरी की आवश्यकता होगी। यह चार A100 GPUs की क्षमता से भी अधिक है। KV cache इतनी तेज़ी से बढ़ता है कि यह अक्सर मॉडल वेट्स (model weights) की तुलना में अधिक मेमोरी का उपयोग करता है।

पारंपरिक मेमोरी मैनेजमेंट विफल हो जाता है क्योंकि:

  • इंटरनल फ्रैगमेंटेशन (Internal fragmentation): आप 4096 टोकन के लिए जगह आवंटित करते हैं लेकिन केवल 300 का उपयोग करते हैं। आप उस जगह का 93% बर्बाद कर देते हैं।
  • शेयरिंग की कमी: एक ही सिस्टम प्रॉम्प्ट वाले दो उपयोगकर्ता उस प्रॉम्प्ट की अपनी अलग-अलग कॉपी स्टोर करते हैं।
  • ऑल-ऑर-नथिंग इविक्शन (All-or-nothing eviction): जब मेमोरी खत्म हो जाती है, तो आपको पूरे सीक्वेंस को CPU पर ले जाना पड़ता है। इससे GPU रुक (stall) जाता है।

PagedAttention इसे कैसे ठीक करता है:

PagedAttention एक ऑपरेटिंग सिस्टम की तरह काम करता है। यह KV cache को छोटे, निश्चित आकार के ब्लॉक्स में विभाजित करता है जिन्हें 'पेज' (pages) कहा जाता है।

यह तीन मुख्य समस्याओं को हल करता है:

  • ऑन-डिमांड एलोकेशन (On-demand allocation): एक सीक्वेंस जैसे-जैसे बढ़ता है, वह केवल उतने ही पेज लेता है। आप अप्रयुक्त क्षमता (unused capacity) पर मेमोरी बर्बाद नहीं करते हैं।
  • शेयर्ड प्रीफ़िक्स सपोर्ट (Shared prefix support): कई उपयोगकर्ता एक सामान्य सिस्टम प्रॉम्प्ट के लिए एक ही फिजिकल पेजों को साझा कर सकते हैं। यह भारी मात्रा में मेमोरी बचाने के लिए "copy-on-write" लॉजिक का उपयोग करता है।
  • फाइन-ग्रेन्ड इविक्शन (Fine-grained eviction): जब मेमोरी भर जाती है, तो सिस्टम बड़े सीक्वेंस के बजाय छोटे पेजों को CPU पर भेज देता है।

परिणाम:

PagedAttention (vLLM के अंदर की तकनीक) का उपयोग करने से पारंपरिक तरीकों की तुलना में थ्रूपुट (throughput) 2x से 4x तक बढ़ सकता है।

इसका उपयोग कब करें:

  • हाई कॉन्करेंसी (High concurrency)।
  • अलग-अलग लंबाई के सीक्वेंस।
  • ऐसे प्रॉम्प्ट जो एक ही शुरुआत साझा करते हैं।

इसे कब छोड़ें (skip करें):

  • सिंगल-यूज़र लोकल इन्फरेंस (Single-user local inference)।
  • बहुत छोटे मॉडल।
  • ऐसे कार्य जहाँ हर सीक्वेंस की लंबाई बिल्कुल समान होती है।

KV Cache और PagedAttention: ये क्या करते हैं और ये क्यों महत्वपूर्ण हैं?

Large Language Models (LLMs) जैसे GPT-4 या Llama-3 के साथ काम करते समय, आपने शायद "Inference" (अनुमान लगाने की प्रक्रिया) के बारे में सुना होगा। लेकिन क्या आपने कभी सोचा है कि ये मॉडल इतने तेज़ कैसे होते हैं, या फिर वे इतनी बड़ी मात्रा में डेटा को कैसे संभालते हैं?

दो महत्वपूर्ण तकनीकें हैं जो इस प्रक्रिया को संभव बनाती हैं: KV Cache और PagedAttention।

इस लेख में, हम समझेंगे कि ये क्या हैं और ये LLM के प्रदर्शन (performance) के लिए क्यों महत्वपूर्ण हैं।

KV Cache क्या है?

LLM के काम करने का तरीका "Autoregressive" होता है। इसका मतलब है कि मॉडल एक बार में एक ही टोकन (token) जनरेट करता है। हर नया टोकन पिछले सभी टोकन पर निर्भर करता है।

मान लीजिए आप मॉडल से एक वाक्य पूरा करने के लिए कह रहे हैं: "आसमान का रंग..."

  1. मॉडल पहले "आसमान" को प्रोसेस करता है।
  2. फिर "का" को प्रोसेस करता है।
  3. फिर "रंग" को प्रोसेस करता है।
  4. अंत में, वह "नीला" जनरेट करता है।

यदि मॉडल हर नए टोकन के लिए पिछले सभी टोकन को फिर से प्रोसेस (re-compute) करे, तो यह बहुत ही धीमा और कंप्यूटेशनल रूप से बहुत महंगा होगा।

KV Cache इसी समस्या का समाधान है।

Transformer आर्किटेक्चर में, 'Attention' मैकेनिज्म के दौरान दो मुख्य चीज़ें बनती हैं: Keys (K) और Values (V)। KV Cache पिछले टोकन के इन 'Key' और 'Value' वेक्टर्स को मेमोरी में स्टोर कर लेता है। जब अगला टोकन जनरेट किया जाता है, तो मॉडल को पिछले टोकन को फिर से प्रोसेस करने की ज़रूरत नहीं होती; वह बस मेमोरी से सीधे KV Cache का उपयोग करता है।

संक्षेप में:

  • बिना KV Cache के: हर नए टोकन के लिए पुराने सभी टोकन को फिर से कैलकुलेट करना पड़ता है (बहुत धीमा)।
  • KV Cache के साथ: पुराने टोकन के परिणाम मेमोरी में सुरक्षित रहते हैं (बहुत तेज़)।

समस्या: मेमोरी का प्रबंधन (The Memory Problem)

हालाँकि KV Cache बहुत उपयोगी है, लेकिन इसके साथ एक बड़ी समस्या जुड़ी है: मेमोरी का अत्यधिक उपयोग और फ्रैगमेंटेशन (Fragmentation)।

जैसे-जैसे कॉन्टेक्स्ट विंडो (context window) बढ़ती है, KV Cache का आकार बहुत बड़ा हो जाता है। पारंपरिक रूप से, सिस्टम को पहले से ही यह अनुमान लगाना पड़ता है कि एक रिक्वेस्ट के लिए कितनी मेमोरी चाहिए। इससे दो मुख्य समस्याएँ पैदा होती हैं:

  1. Internal Fragmentation (आंतरिक फ्रैगमेंटेशन): हम अक्सर ज़रूरत से ज़्यादा मेमोरी एलोकेट कर देते हैं (जैसे कि हम मान लेते हैं कि वाक्य 512 टोकन का होगा, जबकि वह केवल 10 टोकन का है)। बची हुई मेमोरी बेकार पड़ी रहती है।
  2. External Fragmentation (बाहरी फ्रैगमेंटेशन): मेमोरी के टुकड़े अलग-अलग जगहों पर बिखरे होते हैं, जिससे एक बड़ी निरंतर (contiguous) मेमोरी ब्लॉक मिलना मुश्किल हो जाता है।

इन समस्याओं के कारण, GPU की मेमोरी का बहुत सारा हिस्सा बर्बाद हो जाता है, जिससे एक साथ कम रिक्वेस्ट (low throughput) प्रोसेस हो पाती हैं।

PagedAttention क्या है?

यहीं पर PagedAttention की भूमिका आती है। यह तकनीक vLLM लाइब्रेरी के साथ पेश की गई थी और यह ऑपरेटिंग सिस्टम के Virtual Memory और Paging के कॉन्सेप्ट से प्रेरित है।

जिस तरह एक ऑपरेटिंग सिस्टम मेमोरी को छोटे-छोटे "Pages" में विभाजित करता है और उन्हें भौतिक मेमोरी (physical memory) में कहीं भी रख सकता है, PagedAttention भी KV Cache के साथ यही करता है।

यह कैसे काम करता है?

PagedAttention KV Cache को छोटे, निश्चित आकार के Blocks में विभाजित कर देता है।

  • Non-contiguous Storage: इन ब्लॉक्स को GPU मेमोरी में एक साथ (contiguous) रखने की ज़रूरत नहीं है। वे मेमोरी में कहीं भी बिखरे हो सकते हैं।
  • Dynamic Allocation: जैसे-जैसे नए टोकन आते हैं, सिस्टम ज़रूरत के अनुसार नए ब्लॉक्स आवंटित (allocate) करता है।
  • Block Table: एक "Block Table" का उपयोग यह ट्रैक करने के लिए किया जाता है कि कौन सा ब्लॉक किस सीक्वेंस का हिस्सा है।

PagedAttention के लाभ:

  1. शून्य (Near-zero) मेमोरी बर्बादी: क्योंकि हम केवल ज़रूरत के अनुसार ब्लॉक्स आवंटित करते हैं, इसलिए आंतरिक फ्रैगमेंटेशन लगभग समाप्त हो जाता है।
  2. उच्च थ्रूपुट (Higher Throughput): मेमोरी का कुशल उपयोग करने के कारण, हम एक ही समय में बहुत अधिक संख्या में रिक्वेस्ट (batch size) को प्रोसेस कर सकते हैं।
  3. Shared Memory: यह 'Parallel Sampling' जैसी तकनीकों के लिए बहुत उपयोगी है, जहाँ कई रिक्वेस्ट एक ही प्रॉम्प्ट (prompt) का उपयोग करती हैं। PagedAttention एक ही प्रॉम्प्ट के KV Cache ब्लॉक्स को कई रिक्वेस्ट के बीच साझा (share) कर सकता है, जिससे मेमोरी की भारी बचत होती है।

निष्कर्ष

विशेषता पारंपरिक KV Cache PagedAttention
मेमोरी आवंटन निरंतर (Contiguous) ब्लॉक-आधारित (Non-contiguous)
फ्रैगमेंटेशन बहुत अधिक बहुत कम
मेमोरी उपयोग अक्षम (Inefficient) अत्यधिक कुशल (Highly efficient)
थ्रूपुट कम बहुत अधिक

KV Cache ने LLM इन्फरेंस को तेज़ बनाया, लेकिन PagedAttention ने इसे वास्तव में स्केलेबल (scalable) और कुशल बनाया है। यदि आप आज के आधुनिक AI इंफ्रास्ट्रक्चर (जैसे vLLM) का उपयोग कर रहे हैं, तो आप अनजाने में PagedAttention की शक्ति का लाभ उठा रहे हैं।

अधिक AI और मशीन लर्निंग जानकारी के लिए हमारे समुदाय से जुड़ें: https://t.me/GyaanSetuAi