तकनीकी सारांश: PTStore (प्रिफिक्स टेंसर स्टोर)

समस्या विवरण

लार्ज लैंग्वेज मॉडल (LLM) इन्फरेंस वर्कलोड्स हाई-परफॉर्मेंस कंप्यूटिंग (HPC) डेटा सेंटरों में ऊर्जा खपत और संसाधन मांग के मामले में ट्रेनिंग से भी आगे निकलकर प्रमुख लोड बन गए हैं। LLM इन्फरेंस दो चरणों में होता है: प्रीफिल (इनपुट प्रॉम्प्ट को समानांतर रूप से प्रोसेस करना) और डिकोड (टोकन को क्रमिक रूप से जेनरेट करना)। अटेंशन मैकेनिज्म की पुनरावृत्ति गणना (redundant computation) से बचने के लिए, सिस्टम इंटरमीडिएट परिणामों को स्टोर करने के लिए की-वैल्यू (KV) कैश का उपयोग करते हैं।

जबकि अत्याधुनिक रनटाइम्स (जैसे vLLM) एक सिंगल GPU या नोड के भीतर KV कैशिंग को ऑप्टिमाइज़ करते हैं, वे बड़े पैमाने पर महत्वपूर्ण सीमाओं का सामना करते हैं:

1. क्रॉस-नोड पुन: उपयोग का अभाव: मौजूदा सिस्टम अक्सर वितरित कंप्यूट नोड्स के बीच मेमोरी को एकत्रित करने में विफल रहते हैं। यदि एक नोड पर मौजूद रिक्वेस्ट का प्रिफिक्स दूसरे नोड की रिक्वेस्ट के साथ साझा है, तो दूसरा नोड कैश किए गए टेंसरों का पुन: उपयोग करने के बजाय प्रीफिक्स को फिर से कंप्यूट करता है।
2. मेटाडेटा और लेटेंसी बाधाएं: जो दृष्टिकोण वितरित कैशिंग का प्रयास करते हैं (जैसे LMCache, EvoStore), वे अक्सर रिमोट मेमोरी एक्सेस या जटिल मेटाडेटा सिंक्रोनाइज़ेशन (जैसे सिंगल नोड से परे रेडिक्स-अटेंशन को स्केल करना) के कारण उच्च I/O ओवरहेड्स से जूझते हैं।
3. मेमोरी सीमाएं: व्यक्तिगत GPU मेमोरी बड़े कॉन्टेक्स्ट विंडोज़ के लिए अपर्याप्त है, और होस्ट मेमोरी या SSDs पर ऑफलोडिंग करने से होने वाली लेटेंसी कैशिंग के लाभों को समाप्त कर देती है।

मुख्य चुनौती यह है कि कई कंप्यूट नोड्स पर वितरित बड़ी संख्या में GPUs के बीच, अत्यधिक I/O या मेटाडेटा ओवरहेड्स के बिना, KV कैश प्रीफिक्स के स्केलेबल और लो-लेटेंसी पुन: उपयोग को सक्षम करना है।

कार्यप्रणाली: PTStore आर्किटेक्चर

PTStore (प्रिफिक्स टेंसर स्टोर) एक वितरित, रेप्लिकेटेड टेंसर स्टोर है जिसे लोकप्रिय KV कैश प्रीफिक्स को वितरित और रेप्लिकेट करके इन सीमाओं को दूर करने के लिए डिज़ाइन किया गया है। सिस्टम एक क्लाइंट-सर्वर मॉडल का उपयोग करता है जहाँ प्रत्येक कंप्यूट नोड एक सर्वर चलाता है जो स्थानीय और रिमोट GPU क्लाइंट्स को सेवा देने के लिए लोकल होस्ट मेमोरी और SSD को एकत्रित करता है।

मुख्य डिजाइन सिद्धांत

1. इन्क्रीमेंटल टेंसर स्टोरेज (Trie-like संरचना):

   • पूर्ण KV ब्लॉक्स को स्टोर करने के बजाय, PTStore एक नए ऑब्जेक्ट और पहले से स्टोर किए गए ऑब्जेक्ट्स के सबसे लंबे सामान्य प्रिफिक्स (LCP) के बीच इन्क्रीमेंटल अंतर (टेंसर) को स्टोर करता है।
   • यह पूर्ववर्ती दिशाओं में रेडंडेंसी-मुक्त तरीके से प्रीफिक्स को बढ़ने की अनुमति देता है, जो कि एक 'ट्राई' (trie) के समान है, लेकिन इसे टेंसर-लेवल ग्रैनुलैरिटी के माध्यम से लागू किया गया है।
   • कंसोलिडेटेड मेटाडेटा: महंगे डिस्ट्रिब्यूटेड ट्राई ट्रैवर्सल से बचने के लिए, PTStore एक फ्लैट मेटाडेटा संरचना का उपयोग करता है। प्रत्येक ऑब्जेक्ट के मेटाडेटा में यूनिक टेंसर आईडी की एक सूची होती है। एक लोड ऑपरेशन इन आईडी को चेक करने के लिए उन्हें इटरेट करता है कि क्या वे रेप्लिकेशन कैश में स्थानीय रूप से मौजूद हैं; यदि नहीं, तो यह उन्हें "ओनर" सर्वर से रिमोटली फेच करता है।

2. रेप्लिकेशन के साथ डिस्ट्रिब्यूटेड हिरार्किकल कैशिंग:

   • ओन्ड कैश (Owned Cache): उन इन्क्रीमेंटल टेंसरों को स्टोर करता है जिनके लिए एक विशिष्ट सर्वर जिम्मेदार है।
   • रेप्लिकेशन कैश (Replication Cache): एक्सेस लोकैलिटी में सुधार के लिए सर्वर पर "हॉट" (लोकप्रिय) प्रीफिक्स की प्रतियां स्टोर करता है।
   • ट्रेड-ऑफ मैनेजमेंट: सिस्टम ओन्ड और रेप्लिकेटेड कैश के बीच एक कॉन्फ़िगर करने योग्य थ्रेशोल्ड को मैनेज करता है। यह रिट्रीवल स्पीड और स्टोरेज कैपेसिटी के बीच संतुलन बनाने के लिए रेप्लिकेटेड टेंसरों (जिन्हें फिर से फेच किया जा सकता है) को हटाने को प्राथमिकता देता है, बजाय ओन्ड टेंसरों (जिन्हें धीमे स्टोरेज पर फ्लश करने की आवश्यकता होती है) को हटाने के।

3. एक्सेस पैटर्न-अवेयर इविक्शन:

   • PTStore LRU के बजाय एक फ्रीक्वेंसी-बेस्ड इविक्शन पॉलिसी (GDSF से अनुकूलित) का उपयोग करता है, क्योंकि प्रीफिक्स स्ट्रक्चर का अर्थ है कि शुरुआती टेंसरों को अधिक बार एक्सेस किया जाता है।
   • यह साइज बनाम फ्रीक्वेंसी ट्रेड-ऑफ को ध्यान में रखता है, यह सुनिश्चित करता है कि छोटे, बार-बार एक्सेस होने वाले टेंसर बड़े, महंगे-टू-फेच टेंसरों को विस्थापित न करें।

4. RDMA-अवेयर कंसोलिडेशन:

   • बिखराव (scattering) को कम करने के लिए, LCP में जोड़े गए इन्क्रीमेंट्स को ओनर सर्वर पर एक सिंगल कंटीन्यूअस रीजन में कंसोलिडेट किया जाता है।
   • लोड ऑपरेशंस स्कैटर किए गए सेगमेंट को एक सिंगल RPC के माध्यम से समानांतर में फेच करने के लिए बल्क RDMA (रिमोट डायरेक्ट मेमोरी एक्सेस) का उपयोग करते हैं, जिससे ट्रांसफर से पहले डेटा को कंटीन्यूअस रीजन में कॉपी करने का ओवरहेड बच जाता है।

मुख्य योगदान

1. डिजाइन प्रिंसिपल्स: एक वितरित रिपॉजिटरी के लिए सिद्धांतों का एक सेट जो इन्क्रीमेंटल टेंसर स्टोरेज, कंसोलिडेटेड मेटाडेटा और प्रिफिक्स रेप्लिकेशन को एकीकृत करता है।
2. PTStore प्रोटोटाइप: इन सिद्धांतों को लागू करने वाला एक रिसर्च प्रोटोटाइप, जिसमें एक C++ लो-लेवल API और vLLM जैसे LLM रनटाइम्स के साथ सहज एकीकरण के लिए एक पायथन इंटरफेस है।
3. परफॉर्मेंस वैलिडेशन: अत्याधुनिक बेसलाइन्स की तुलना में I/O ओवरहेड और एंड-टू-एंड रनटाइम में महत्वपूर्ण कमी प्रदर्शित करने वाले व्यापक प्रयोग।

प्रयोगात्मक परिणाम

लेखकों ने दो एक्सट्रैक्टिव QA वर्कलोड्स: WikiQA (लॉन्ग कॉन्टेक्स्ट) और SQUAD (हाई वॉल्यूम ऑफ क्वेश्चंस) का उपयोग करके ALCF Polaris HPC टेस्टबेड (560 नोड्स, A100 GPUs) पर PTStore का मूल्यांकन किया। उपयोग किया गया LLM Mistral-7B-instruct-V2 था।

बेसलाइन्स

• vLLM Vanilla: बिना किसी क्रॉस-रिक्वेस्ट प्रिफिक्स शेयरिंग के स्टैंडर्ड vLLM।
• vLLM Prefix: लोकल प्रिफिक्स शेयरिंग (एक नोड के भीतर) के साथ vLLm।
• EvoStore: इन्क्रीमेंटल स्टोरेज और RDMA का उपयोग करने वाला एक डिस्ट्रिब्यूटेड टेंसर स्टोर, लेकिन इसमें लोकल प्रिफिक्स रेप्लिकेशन की कमी है।
• PTStore: प्रस्तावित सिस्टम जिसमें डिस्ट्रिब्यूटेड अवेयरनेस और लोकल रेप्लिकेशन है।

निष्कर्ष

• कम स्केलेबिलिटी (8–32 GPUs): PTStore ने EvoStore और vLLM Prefix को काफी पीछे छोड़ दिया। जबकि EvoStore रिमोट प्रीफिक्स को फेच करने के दौरान उच्च RDMA I/O ओवरहेड्स से जूझ रहा था, PTStore के लोकल रेप्लिकेशन ने इसे कम किया, जिसके परिणामस्वरूप टाइम टू फर्स्ट टोकन (TTFT) में एक "डिटैच्ड एडवांटेज" मिला।
• सीक्वेंस लेंथ स्केलेबिलिटी (1k–8k टोकन):
  • शॉर्ट सीक्वेंस (1k) के लिए, vLLM की लोकल कैशिंग प्रतिस्पर्धी थी।
  • जैसे-जैसे सीक्वेंस लेंथ बढ़ी, PTStore का लाभ बढ़ता गया। 8k टोकन पर, PTStore vLLM के प्रिफिक्स कैशिंग से लगभग 2 गुना तेज़ और EvoStore से 20% तेज़ था।
  • लंबे कॉन्टेक्स्ट के साथ प्रदर्शन का अंतर बढ़ गया क्योंकि रीकंप्यूटेशन या रिमोट I/O की लागत लोकल-ओनली कैशिंग के लाभों से अधिक हो गई।
• एफिशिएंसी गेन्स: लंबे पैसेज Q&A डेटासेट्स पर, PTStore ने उन बेसलाइन्स की तुलना में 5–6 गुना अधिक कुशलता से इन्फरेंस निष्पादित किया जो नोड्स के बीच मेमोरी को एकत्रित नहीं करते हैं और जिन्हें KV कैश को फिर से जेनरेट करने की आवश्यकता होती है।

महत्व और दावे

पेपर का दावा है कि PTStore स्केलेबल LLM इन्फरेंस सर्विंग के एक महत्वपूर्ण अंतर को संबोधित करता है: वितरित नोड्स के बीच KV कैश प्रीफिक्स को कुशलतापूर्वक पुन: उपयोग करने में वर्तमान सिस्टमों की अक्षमता। इन्क्रीमेंटल स्टोरेज को न्यूनतम रेडंडेंसी, कंसोलिडेटेड मेटाडेटा को त्वरित क्वेरीज़ के लिए, और रेप्लिकेशन रणनीति को लोकैलिटी को ऑप्टिमाइज़ करने के लिए जोड़कर, PTStore सक्षम बनाता है:

• क्लस्टर में मेमोरी को एकत्रित करके प्रभावी KV कैश साइज का क्रमों के परिमाण (orders of magnitude) तक विस्तार।
• TTFT में महत्वपूर्ण कमी, विशेष रूप से लंबे कॉन्टेक्स्ट वाले वर्कलोड के लिए जहाँ रीकंप्यूटेशन महंगा होता है।
• स्केलेबिलिटी जो पिछले वितरित दृष्टिकोणों को परेशान करने वाले संचार बॉटलनैक्स और मेटाडेटा सिंक्रोनाइज़ेशन समस्याओं से बचती है।

लेखक PTStore को स्केलेबल AI इन्फरेंस की ओर एक मौलिक कदम के रूप में देखते हैं, यह नोट करते हुए कि भविष्य का कार्य डायनेमिक मेमोरी बैलेंसिंग, ML-आधारित इविक्शन पॉलिसीज़ और वास्तविक दुनिया के कन्वर्सेशनल और कोड-कंप्लीशन ट्रेसेस पर LMCache और Mooncake जैसे सिस्टम के विरुद्ध व्यापक बेंचमार्किंग पर केंद्रित होगा।