🤖 AI

OxyGen: Unified KV Cache Management for Vision-Language-Action Models under Multi-Task Parallelism

OxyGen विज़न-लैंग्वेज-एक्शन मॉडल्स के लिए एक एकीकृत KV कैश प्रबंधन प्रतिमान है जो KV कैश को एक साझा संसाधन के रूप में मानकर ऑन-डिवाइस मल्टी-टास्क पैरेललिज्म को कुशल बनाता है, जिससे अनावश्यक गणना समाप्त होती है और गुणवत्ता में गिरावट के बिना भाषा और एक्शन जनरेशन दोनों में महत्वपूर्ण गति प्राप्त होती है।

मूल लेखक: Xiangyu Li, Huaizhi Tang, Xin Ding, Weijun Wang, Ting Cao, Yunxin Liu

प्रकाशित 2026-03-17

📖 5 मिनट में पढ़ें🧠 गहराई से पढ़ें

CC BY 4.0

मूल लेखक: Xiangyu Li, Huaizhi Tang, Xin Ding, Weijun Wang, Ting Cao, Yunxin Liu

मूल पेपर CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) के तहत लाइसेंस किया गया है। ✨ नीचे दिए गए पेपर की यह व्याख्या AI से तैयार की गई है। इसे लेखकों ने न तो लिखा है, न इसका समर्थन किया है। तकनीकी सटीकता के लिए मूल पेपर देखें। पूरा डिस्क्लेमर पढ़ें

कल्पना कीजिए कि आपके पास "रोबो" (Robo) नाम का एक अत्यधिक बुद्धिमान रोबोट सहायक है। रोबो एक साथ तीन काम करने की कोशिश कर रहा है:

कप उठाने के लिए अपना हाथ हिलाना (Action)।
वह जो कर रहा है उसके बारे में आपसे बात करना (Language)।
बाद के लिए उसने कप कहाँ रखा, इसे याद रखना (Memory)।

अतीत में, यदि आप रोबो से ये सभी काम करने के लिए कहते, तो यह ऐसा होता जैसे आपने एक ही काम के लिए तीन अलग-अलग लोगों को काम पर रखा हो, लेकिन वे सभी एक छोटे से किचन में एक ही चूल्हे का उपयोग करने की कोशिश कर रहे हों। वे आपस में टकराते रहते, लाइन में प्रतीक्षा करते और बार-बार एक ही सामग्री को दोबारा पकाते रहते। परिणाम? रोबो धीमा, अनाड़ी हो जाता था और अक्सर कप गिरा देता था क्योंकि बात करने में व्यस्त होने के कारण वह अपने हाथ को पर्याप्त तेज़ी से नहीं हिला पाता था।

यह पेपर OxyGen को पेश करता है, जो रोबो के मस्तिष्क को प्रबंधित करने का एक नया तरीका है ताकि वह सुचारू रूप से मल्टीटास्किंग कर सके। यह कैसे काम करता है, इसके लिए सरल उपमाओं का उपयोग किया गया है:

समस्या: "अलग-थलग किचन" (The Isolated Kitchen)

वर्तमान रोबोट सिस्टम हर कार्य को एक अलग, अलग-थलग किचन की तरह मानते हैं।

दोहराव वाली कुकिंग (Redundant Cooking): यदि रोबो को हाथ हिलाने के लिए और कप के बारे में बात करने के लिए एक तस्वीर देखने की आवश्यकता है, तो पुराना सिस्टम रोबोट को उस तस्वीर को दो बार "पकाने" (प्रोसेस करने) के लिए मजबूर करता है। यह समय और ऊर्जा की बर्बादी है।
ट्रैफिक जाम (The Traffic Jam): भले ही वे एक ही चूल्हा साझा कर रहे हों, "बात करने" वाले कार्य को एक वाक्य पूरा करने में लंबा समय लग सकता है। जब तक वह बात कर रहा होता है, "हिलाने" वाले कार्य को प्रतीक्षा करनी पड़ती है, भले ही हाथ हिलाने के लिए तुरंत प्रतिक्रिया (रिफ्लेक्स) की आवश्यकता हो। इससे रोबोट रुक जाता है या झटकेदार गति करने लगता है।

समाधान: OxyGen (एक एकीकृत प्रबंधक - The Unified Manager)

OxyGen एक अत्यधिक कुशल हेड शेफ की तरह कार्य करता है जो पूरे किचन को एक बड़े, साझा स्थान के रूप में प्रबंधित करता है। यह दो मुख्य तरकीबों का उपयोग करता है:

1. "साझा रेसिपी बुक" (The Shared Recipe Book - Cross-Task KV Sharing)

AI में, "KV कैश" एक रफ नोटपैड या वर्किंग मेमोरी की तरह है जहाँ रोबोट वह लिखता है जो उसने पहले देखा और समझा है।

पुराना तरीका: जब भी रोबට कोई नया कार्य शुरू करता है (हिलाना, बात करना, याद रखना), तो वह पुराने रफ नोटपैड को फेंक देता है और शून्य से लिखना शुरू कर देता है, भले ही वह उसी कप को देख रहा हो।
OxyGen का तरीका: हेड शेफ कहता है, "हे, हमने पहले ही कप को देख लिया है! चलो उस पेज को रफ नोटपैड में रहने देते हैं।"
परिणाम: रोबोट को दोबारा कप को पढ़ने में समय बर्बाद नहीं करना पड़ता। यह तुरंत उस समझ को हाथ हिलाने वाले कार्य और बात करने वाले कार्य के बीच साझा कर देता है। यह तीन अलग-अलग वर्ड फाइलों को आपस में ईमेल करने के बजाय एक सिंगल गूगल डॉक साझा करने जैसा है।

2. "असेंबली लाइन" (The Assembly Line - Cross-Frame Continuous Batching)

रोबोटों की डेडलाइन अलग-अलग होती है। हाथ हिलाने के लिए अभी (हर 1/60 सेकंड में) काम करना ज़रूरी है। बात करना थोड़ा धीमा हो सकता है, जो कई सेकंडों तक चलता है।

पुराना तरीका: रोबोट अपने हाथ को फिर से हिलाने से पहले एक पूरा वाक्य खत्म करने के लिए सब कुछ रोक देता है। इससे हाथ की गति रुक-रुक कर या झटकेदार हो जाती है।
OxyGen का तरीका: एक असेंबली लाइन की कल्पना करें। जब रोबोट इस सेकंड के लिए अपना हाथ हिला रहा होता है, तब वह बैकग्राउंड में वाक्य के अगले कुछ शब्दों पर भी चुपचाप काम कर रहा होता है। यह बात करने वाले कार्यों को एक साथ समूहबद्ध (batch) करता है ताकि कंप्यूटर चिप एक साथ उन सभी पर कुशलतापूर्वक काम कर सके।
परिणाम: हाथ सुचारू रूप से और तेज़ी से चलता है (इंसान की तरह), जबकि रोबोट बिना धीमा हुए साथ ही साथ एक लंबी, विस्तृत कहानी भी सुना रहा होता है।

वास्तविक दुनिया का प्रभाव

शोधकर्ताओं ने इसे एक शक्तिशाली रोबोट ब्रेन (जिसे $\pi0.5$ कहा जाता है) पर एक मानक गेमिंग ग्राफिक्स कार्ड (RTX 4090) के साथ टेस्ट किया।

गति (Speed): रोबोट 3.7 गुना तेज़ हो गया।
सुचारूता (Smoothness): यह 70 बार प्रति सेकंड (अत्यधिक स्मूथ) की दर से अपना हाथ हिला सका, जबकि साथ ही बहुत तेज़ गति (200 शब्द प्रति सेकंड) से बात भी कर रहा था।
गुणवत्ता (Quality): महत्वपूर्ण बात यह है कि इसने गलतियाँ नहीं कीं। रोबोट ने अभी भी कप को सही ढंग से उठाया; इसने बस इसे बहुत तेज़ी से और ट्रैफिक जाम के कारण "थके बिना" किया।

सारांश

OxyGen को एक ऐसे अराजक किचन और एक सुव्यवस्थित किचन के बीच के अंतर के रूप में समझें जहाँ तीन रसोइये एक ही चूल्हे के लिए लड़ रहे हों, बनाम एक सुव्यवस्थित किचन जिसमें एक हेड शेफ है जो सामग्री साझा करता है और एक असेंबली लाइन चलाता है। यह हमारे रोबोट दोस्तों को अंततः वह करने की अनुमति देता है जो इंसान सबसे अच्छा करते हैं: चलना और च्युइंग गम चबाना एक साथ, लेकिन एक मशीन की गति और सटीकता के साथ।

यहाँ शोध पत्र "OxyGen: Multi-Task Parallelism के तहत Vision-Language-Action Models के लिए Unified KV Cache Management" का विस्तृत तकनीकी सारांश दिया गया है।

1. समस्या विवरण (Problem Statement)

एम्बोडेड एआई (Embodied AI) एजेंटों को तेजी से मल्टी-टास्क पैरेललिज्म (multi-task parallelism) की आवश्यकता हो रही है, जहाँ एक एकल एजेंट को साझा विजुअल ऑब्जर्वेशन के आधार पर विविध कार्यों (जैसे, रोबोटिक मैनिपुलेशन, संवादात्मक नरेशन और मेमोरी कंस्ट्रक्शन) को एक साथ निष्पादित करना होता है। हालाँकि हालिया Mixture-of-Transformers (MoT) Vision-Language-Action (VLA) मॉडल (जैसे $\pi^0.5$ ) आर्किटेक्चरल रूप से एक ही मॉडल से हेट्रोजेनियस आउटपुट (एक्शन और लैंग्वेज) उत्पन्न करने का समर्थन करते हैं, मौजूदा इन्फरेंस सिस्टम संसाधन-सीमित उपकरणों (जैसे, सिंगल GPU) पर इसका कुशलतापूर्वक लाभ उठाने में विफल रहते हैं।

यह शोध पत्र आइसोलेटेड KV कैश मैनेजमेंट (isolated KV cache management) को वर्तमान प्रणालियों में अक्षमता के मूल कारण के रूप में पहचानता है। मौजूदा दृष्टिकोण प्रत्येक कार्य को स्वतंत्र रूप से देखते हैं, जिससे दो महत्वपूर्ण बाधाएं उत्पन्न होती हैं:

रिडंडेंट कंप्यूटेशन (Redundant Computation): साझा ऑब्जर्वेशन को प्रत्येक कार्य के लिए बार-बार एनकोड किया जाता है, जिससे समान Key-Value (KV) कैश प्रविष्टियाँ कई बार उत्पन्न होती हैं।
रिसोर्स कंटेंशन (Resource Contention): कार्य बिना किसी समन्वय के सीमित हार्डवेयर संसाधनों (GPU मेमोरी/बैंडविड्थ) के लिए प्रतिस्पर्धा करते हैं। महत्वपूर्ण रूप से, कार्यों की असिमेट्रिक डेडलाइन्स (asymmetric deadlines) होती हैं: स्मूथ कंट्रोल के लिए एक्शन जनरेशन को हार्ड, रियल-टाइम डेडलाइन (जैसे, 50–70 Hz) की आवश्यकता होती है, जबकि लैंग्वेज जनरेशन में सॉफ्ट डेडलाइन होती है और यह कई फ्रेम्स तक चल सकता है। वर्तमान सिस्टम इन बाधाओं को अलग करने में सक्षम नहीं हैं, जिससे महत्वपूर्ण धीमापन (2.6 $\times$ तक) और कम थ्रूपुट होता है।

2. कार्यप्रणाली: OxyGen

लेखक OxyGen का प्रस्ताव करते हैं, जो एक नवीन प्रतिमान (paradigm) पर आधारित एक इन्फरेंस सिस्टम है जिसे Unified KV Cache Management कहा जाता है। यह दृष्टिकोण KV कैश को कार्य-विशिष्ट आर्टिफैक्ट के बजाय, कार्यों और समय के बीच एक फर्स्ट-क्लास, साझा संसाधन के रूप में मानता है।

मुख्य घटक:

Unified KV Cache Manager ( $\mathcal{M}$ ): एक केंद्रीय शेड्यूलर जो सभी इन-फ्लाईट अनुरोधों की स्थिति बनाए रखता है। यह सभी इन-फ्लाईट अनुरोधों की स्थिति को ट्रैक करता है जिसमें KV कैश, टोकन बफ़र्स और टर्मिनेशन फ्लैग्स शामिल हैं।
ऑप्टिमाइज़ेशन 1: क्रॉस-टास्क KV शेयरिंग (Cross-Task KV Sharing):
- एक्शन एक्सपर्ट और लैंग्वेज एक्सपर्ट दोनों के लिए विजुअल ऑब्जर्वेशन को दोबारा एनकोड करने के बजाय, सिस्टम एक सिंगल प्रीफिल पास (single prefill pass) करता है ताकि साझा KV कैश ( $K_t$ ) उत्पन्न किया जा सके।
- इसके बाद इस $K_t$ को उसी फ्रेम के भीतर दोनों एक्सपर्ट्स द्वारा डुप्लिकेट और कंज्यूम किया जाता है, जिससे रिडंडेंट प्रीफिल कंप्यूटेशन समाप्त हो जाता है।
ऑप्टिमाइज़ेशन 2: क्रॉस-फ्रेम कंटीन्यूअस बैचिंग (Cross-Frame Continuous Batching):
- सिस्टम लैंग्वेज डिकोडिंग को कठोर प्रति-फ्रेम कंट्रोल लूप से अलग करता है।
- एक्शन टास्क को वर्तमान फ्रेम के भीतर तुरंत प्रोसेस किया जाता है ताकि हार्ड डेडलाइन को पूरा किया जा सके।
- लैंग्वेज टास्क को स्ट्रीमिंग अनुरोधों के रूप में माना जाता है। मैनेजर विभिन्न फ्रेम्स से कई सक्रिय लैंग्वेज अनुरोधों को एक एकल बैच स्टेट ( $\hat{\sigma}$ ) में एकत्रित करता है।
- VLM इस बैच पर समानांतर में ऑटोरिग्रेसिव डिकोडिंग करता है, जिससे कई अनुरोधों के बीच डिकोडिंग लागत का बंटवारा (amortize) होता है और GPU यूटिलाइजेशन अधिकतम होता है।

निष्पादन प्रवाह (प्रति फ्रेम - Execution Flow):

प्रीफिल (Prefill): साझा $K_t$ उत्पन्न करने के लिए नए ऑब्जर्वेशन $o_t$ को एक बार एनकोड करें।
एक्शन जनरेशन (Action Generation): एक्शन चंक $A_t$ (डिफ्यूजन/फ्लो मैचिंग के माध्यम से) उत्पन्न करने के लिए तुरंत $K_t$ का उपयोग करें।
स्टेट मैनेजमेंट (State Management): $K_t$ का उपयोग करके एक नया लैंग्वेज रिक्वेस्ट स्टेट $\sigma_t$ शुरू करें और इसे मैनेजर में स्टोर करें।
बैच्ड डिकोडिंग (Batched Decoding): सभी सक्रिय लैंग्वेज अनुरोधों (नए वाले सहित) को प्राप्त करें, उन्हें बैच करें, और $k$ स्टेप्स की पैरेलल डिकोडिंग करें।
अपडेट (Update): परिणामों को अनबैच करें, नए टोकन्स के साथ मैनेजर को अपडेट करें, और पूर्ण अनुरोधों को हटा दें (evict)।

3. मुख्य योगदान

समस्या का निरूपण (Problem Formulation): यह शोध पत्र असिमेट्रिक डेडलाइन के साथ मल्टी-टास्क पैरेललिज्म को MoT-VLAs के लिए एक महत्वपूर्ण इन्फरेंस परिदृश्य के रूप में औपचारिक रूप से परिभाषित करता है और आइसोलेटेड KV कैश मैनेजमेंट को प्राथमिक बाधा के रूप में पहचानता है।
Unified KV Cache प्रतिमान: एक नया इन्फरेंस प्रतिमान पेश करता है जो KV कैश को एक साझा संसाधन के रूप में एब्स्ट्रैक्ट करता है, जिससे क्रॉस-टास्क शेयरिंग (रिडंडेंट प्रीफिल को हटाना) और क्रॉस-फ्रेम कंटीन्यूअस बैचिंग (वेरिएबल-लेंथ लैंग्वेज डिकोडिंग को फिक्स्ड-रेट एक्शन जनरेशन से अलग करना) संभव होता है।
सिस्टम कार्यान्वयन और मूल्यांकन: $\pi^0.5$ (सबसे लोकप्रिय MoT VLA) के लिए OxyGen को लागू किया गया और NVIDIA RTX 4090 का उपयोग करके प्रतिनिधि रोबोटिक बेंचमार्क (LIBERO, DROID, ALOHA) पर इसका सत्यापन किया गया।

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

एक सिंगल NVIDIA RTX 4090 GPU पर मूल्यांकित, OxyGen दोनों सीक्वेंशियल और नेइव पैरेलल बेसलाइन्स (CUDA MPS का उपयोग करके) की तुलना में महत्वपूर्ण प्रदर्शन लाभ प्रदर्शित करता है:

स्पीडअप (Speedup): आइसोलेटेड निष्पादन की तुलना में एक्शन फ्रीक्वेंसी और लैंग्वेज थ्रूपुट में 3.7 $\times$ तक की स्पीडअप प्राप्त करता है।
थ्रूपुट (Throughput): कंट्रोल क्वालिटी को कम किए बिना लैंग्वेज डिकोडिंग के लिए >200 टोकन/सेकंड और एक्शन जनरेशन के लिए 70 Hz को एक साथ प्राप्त करता है।
एब्लेशन स्टडीज (Ablation Studies):
- केवल क्रॉस-टास्क KV शेयरिंग रिडंडेंट प्रीफिल को हटाकर ~1.4 $\times$ स्पीडअप प्रदान करती है।
- क्रॉस-फ्रेम कंटीन्यूअस बैचिंग लंबे डिकोडिंग कार्यों के लिए महत्वपूर्ण है, जो हाई एक्शन फ्रीक्वेंसी (लगभग 60 Hz) को बनाए रखती है, जबकि बेसलाइन्स ~19 Hz तक गिर जाते हैं।
क्वालिटी और एफिशिएंसी:
- एक्शन क्वालिटी: LIBERO बेंचमार्क पर टास्क सक्सेस रेट मूल $\pi^0.5$ के प्रदर्शन के समान है (सांख्यिकीय शोर के भीतर), जो पुष्टि करता है कि कंट्रोल क्वालिटी में कोई गिरावट नहीं आई है।
- ऊर्जा दक्षता: बेसलाइन्स की तुलना में प्रति अनुरोध ऊर्जा में 78% तक की कमी और औसत बिजली खपत में 47% की कमी आती है, जिसका कारण कम मेमोरी एक्सेस और बेहतर हार्डवेयर यूटिलाइजेशन है।
- मेमोरी: बेसलाइन की तुलना में केवल ~15% मेमोरी ओवरहेड जोड़ता है, जबकि नेइव पैरेललाइजेशन लगभग दोगुना मेमोरी उपयोग करता है।

5. महत्व

OxyGen उन्नत एम्बोडेड एआई एजेंटों को तैनात करने में एक महत्वपूर्ण अंतर को पाटता है। KV कैश मैनेजमेंट को एक यूनिफाइड रिसोर्स के रूप में फिर से सोचकर, यह MoT-VLAs को उनके आर्किटेक्चरल पोटेंशियल को सिमल्टेनियस मल्टी-टास्क एग्जीक्यूशन के लिए वास्तव में साकार करने में सक्षम बनाता है। यह कार्य प्रदर्शित करता है कि कंट्रोल फ्रीक्वेंसी या लैंग्वेज फ्लूएंसी से समझौता किए बिना जटिल, मल्टी-मोडल एजेंटों के लिए कुशल ऑन-डिवाइस इन्फरेंस संभव है, जो अधिक सक्षम, रियल-टाइम और ऊर्जा-कुशल होम और इंडस्ट्रियल रोबोट के लिए मार्ग प्रशस्त करता है।