🤖 machine learning

StreamKL: Fast and Memory-Efficient KL Divergence for Boosting Attention Distillation

StreamKL एक नवीन फ्यूज्ड GPU प्रिमिटिव (fused GPU primitive) पेश करता है जो क्वेरी-की टाइल्स (query-key tiles) को एक ही पास में स्ट्रीम करके अटेंशन डिस्टिलेशन (attention distillation) की द्विघातीय मेमोरी और I/O बाधाओं को समाप्त करता है, जिससे एक सिंगल GPU पर लॉन्ग-कॉन्टेक्स्ट डिस्टिलेशन सक्षम करने के लिए मेमोरी फुटप्रिंट को $O(N_QN_K)$ से घटाकर $O(1)$ करने के साथ महत्वपूर्ण गति प्राप्त होती है।

मूल लेखक: Guangda Liu, Yiquan Wang, Chengwei Li, Wenhao Chen, Jing Lin, Yiwu Yao, Danning Ke, Wenchao Ding, Jieru Zhao

प्रकाशित 2026-06-19

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

CC BY 4.0

मूल लेखक: Guangda Liu, Yiquan Wang, Chengwei Li, Wenhao Chen, Jing Lin, Yiwu Yao, Danning Ke, Wenchao Ding, Jieru Zhao

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

कल्पना कीजिए कि आप एक छोटे, तेज़ छात्र (एक "स्टूडेंट मॉडल") को एक बुद्धिमान, धीमे शिक्षक (एक "टीचर मॉडल") की तरह सोचना सिखाने की कोशिश कर रहे हैं। AI की दुनिया में, वे यह देखने के लिए कि दोनों किसी कहानी या वाक्य के विभिन्न हिस्सों पर कैसे "ध्यान देते" (pay attention) हैं, उनकी तुलना करते हैं। इस प्रक्रिया को अटेंशन डिस्टिलेशन (Attention Distillation) कहा जाता है।

इस तुलना को करने के लिए, कंप्यूटर एक विशिष्ट संख्या की गणना करता है जिसे KL डाइवर्जेंस (KL Divergence) कहा जाता है। इसे एक "दूरी स्कोर" (distance score) के रूप में समझें जो आपको बताता है कि छात्र का ध्यान शिक्षक के ध्यान से कितना अलग है। लक्ष्य इस स्कोर को यथासंभव छोटा करना है।

समस्या: "मेमोरी विस्फोट" (The "Memory Explosion")

पेपर बताता है कि लंबी कहानियों (जैसे 100,000 शब्दों वाला उपन्यास) के लिए यह तुलना करना वर्तमान में कंप्यूटर मेमोरी के लिए एक दुस्वप्न है।

यहाँ इसकी उपमा (analogy) दी गई है:
कल्पना कीजिए कि आपके पास दो विशाल व्हाइटबोर्ड हैं, एक शिक्षक के अटेंशन के लिए और एक छात्र के अटेंशन के लिए। उनकी तुलना करने के लिए, पुराने तरीके में आपको शब्दों के हर एक संभावित जोड़े (pairing) को इन बोर्डों पर लिखना पड़ता है।

यदि आपके पास 64,000 शब्द हैं, तो आपको 64,000 × 64,000 जोड़े लिखने होंगे। यह 4 अरब से अधिक संख्याएँ हैं।
ऐसा करने के लिए एक ऐसे विशाल व्हाइटबोर्ड की आवश्यकता होती है जो कंप्यूटर की मुख्य मेमोरी (HBM) में फिट न हो सके। यह एक जूते के डिब्बे में किताबों के पुस्तकालय को रखने की कोशिश करने जैसा है।
क्योंकि कंप्यूटर एक बार में पूरी तस्वीर नहीं देख सकता, इसलिए उसे कहानी को छोटे-छोटे टुकड़ों में काटना पड़ता है, उन्हें प्रोसेस करना पड़ता है, और फिर उन्हें वापस जोड़ना पड़ता है। यह धीमा है, जैसे एक किताब को एक समय में एक अक्षर देखकर पढ़ने की कोशिश करना और अगले पर जाने से पहले उसे लिखना।

समाधान: StreamKL (एक "स्ट्रीमिंग" दृष्टिकोण)

लेखकों ने StreamKL नामक एक नया टूल बनाया है। सब कुछ एक विशाल व्हाइटबोर्ड पर लिखने के बजाय, StreamKL चतुराई से "डिस्टेंस स्कोर" की गणना ऑन-द-फ्लाई (on the fly) करता है, जैसे कि एक कन्वेयर बेल्ट हो।

रचनात्मक उपमा: एक फैक्ट्री असेंबली लाइन
कल्पना कीजिए कि एक फैक्ट्री है जहाँ आप उत्पादों के दो कन्वेयर बेल्ट (शिक्षक का अटेंशन और छात्र का अटेंशन) की तुलना कर रहे हैं।

पुराना तरीका: आप लाइन रोक देते हैं, हर उत्पाद को एक विशाल गोदाम के फर्श (HBM) पर डाल देते हैं, उन सभी को मापते हैं, और फिर सफाई करते हैं। इसमें पूरे गोदाम की जगह लगती है और यह धीमा है।
StreamKL का तरीका: आप उत्पादों को कन्वेयर बेल्ट पर चलते रहने देते हैं। जैसे-जैसे प्रत्येक जोड़ी आइटम एक सेंसर (GPU चिप) के पास से गुजरता है, आप तुरंत उनकी तुलना करते हैं, अंतर की गणना करते हैं, और अगले जोड़े के आने से पहले परिणाम को एक छोटी जेब (SRAM) में डाल देते हैं। आपको लाइन कभी नहीं रोकनी पड़ती, और आपको कभी भी गोदाम की आवश्यकता नहीं होती। आपको केवल एक जेब की आवश्यकता होती है।

यह कैसे काम करता है (जादुई ट्रिक)

पेपर इस जादू के दो मुख्य भागों का वर्णन करता है:

फॉरवर्ड पास (स्कोर की गणना करना):
शोधकर्ताओं ने एक नया गणितीय सूत्र आविष्कार किया है जो कंप्यूटर को डेटा के माध्यम से धीरे-धीरे (incrementally) "डिस्टेंस स्कोर" को अपडेट करने की अनुमति देता है। जैसे-जैसे यह डेटा के माध्यम से स्ट्रीम होता है, यह पूरी सूची के बजाय केवल कुछ संख्याओं (जैसे कि एक रनिंग मैक्सिमम और एक सम/योग) का हिसाब रखता है। इसका मतलब है कि यह बिना मेमोरी खत्म किए किसी भी लंबाई की कहानियों को संभाल सकता है।
बैकवर्ड पास (गलतियों से सीखना):
जब कंप्यूटर को छात्र को सुधारने के लिए स्कोर से सीखने की आवश्यकता होती है, तो उसे आमतौर पर डेटा को पीछे देखना पड़ता है। पुराना तरीका पीछे देखने के लिए पूरी विशाल सूची को सहेज लेता है। StreamKL स्मार्ट है: यह सूची को फेंक देता है लेकिन कुछ "गुप्त चाबियाँ" (जिन्हें LSE मान कहा जाता है) याद रखता है। जब इसे पीछे देखने की आवश्यकता होती है, तो यह इन चाबियों का उपयोग करके उस विशिष्ट डेटा को पुनर्निर्मित (rebuild) करता है जिसकी इसे आवश्यकता है, फिर उसी समय सबक की गणना करता है, और फिर इसे भूल जाता है। यह केक की रेसिपी याद रखने जैसा है ताकि जब भी आपको एक टुकड़ा चखने की आवश्यकता हो, तो आप उसे बना सकें, बजाय इसके कि पूरा केक बनाकर फ्रिज में रख दें।

परिणाम: गति और स्थान

पेपर ने बहुत लंबे संदर्भों (512,000 शब्दों तक) के साथ शक्तिशाली NVIDIA GPUs (H200 और A100) पर इसका परीक्षण किया।

मेमोरी बचत: StreamKL ने अतिरिक्त मेमोरी की आवश्यकता को "क्वाड्रेटिक" (टेराबाइट्स तक विस्फोट होना) से घटाकर "कॉन्स्टेंट" (बहुत कम रहना) कर दिया। इसे 64k कॉन्टेक्स्ट के लिए 512 GB मेमोरी की आवश्यकता से बदलकर लगभग कुछ भी अतिरिक्त आवश्यकता नहीं रह गई। यह एक सिंगल GPU को उन कार्यों को संभालने की अनुमति देता है जिनके लिए पहले एक सुपरकंप्यूटर या जो असंभव थे, उनकी आवश्यकता थी।
गति: क्योंकि इसे डेटा के बड़े पैमाने पर आदान-प्रदान करने के लिए लिखने और पढ़ने की आवश्यकता नहीं होती है, इसलिए यह अविश्वसनीय रूप से तेज़ है।
- कुछ परीक्षणों में, यह स्कोर की गणना करने के लिए मानक विधि की तुलना में 43 गुना तेज़ था।
- सीखने के चरण में, यह 14 गुना तेज़ था।

सारांश

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

तकनीकी सारांश: StreamKL

समस्या विवरण

अटेंशन डिस्टिलेशन (Attention distillation), जो एक अटेंशन डिस्ट्रीब्यूशन को दूसरे से मेल खाने के लिए कुलबैक-लीब्लर (KL) डाइवर्जेंस को कम करने की एक तकनीक है, नॉलेज डिस्टिलेशन, मॉडल प्रूनिंग, कॉन्टिनुअल लर्निंग और स्पार्स-अटेंशन लार्ज लैंग्वेज मॉडल्स (LLMs) के प्रशिक्षण के लिए महत्वपूर्ण है। हालांकि, मौजूदा कार्यान्वयन लंबे संदर्भ लंबाई ( $N_Q, N_K$ ) को संभालने के दौरान गंभीर स्केलेबिलिटी बाधाओं का सामना करते हैं।

वर्तमान दृष्टिकोण दोनों अटेंशन डिस्ट्रीब्यूशन ( $P_1, P_2 \in \mathbb{R}^{N_Q \times N_K}$ ) को एलिमेंट-वाइज KL रिडक्शन कंप्यूट करने से पहले हाई बैंडविड्थ मेमोरी (HBM) में "मटेरियलाइज़" (materialize) करते हैं। इसके परिणामस्वरूप:

क्वाड्रेटिक मेमोरी फुटप्रिंट (Quadratic Memory Footprint): दो $N_Q \times N_K$ मैट्रिसेस को स्टोर करने के लिए $O(N_Q N_K)$ मेमोरी की आवश्यकता होती है। 32 हेड्स के साथ 64K टोकन के संदर्भ के लिए, BF16 में यह अकेले ~512 GB उपभोग करता है, जो एक एकल NVIDIA H200 GPU की 141 GB क्षमता से अधिक है।
क्वाड्रेटिक IO लागत (Quadratic IO Costs): इन डिस्ट्रीब्यूशन को HBM में लिखना और उन्हें KL रिडक्शन के लिए वापस पढ़ना, संदर्भ लंबाई के साथ क्वाड्रेटिक रूप से स्केल होता है और रनटाइम पर हावी हो जाता है।
चंकिंग की अक्षमता (Inefficiency of Chunking): सामान्य वर्कअराउंड जैसे चंक्ड प्रोसेसिंग (chunked processing) मेमोरी के बदले लेटेंसी का व्यापार करते हैं; छोटे चंक्स GPU पैरेललिज्म का कम उपयोग करते हैं, जबकि बड़े चंक्स आउट-ऑफ-मेमोरी (OOM) त्रुटियों को ट्रिगर करते हैं।

स्टैंडर्ड अटेंशन कंप्यूटेशन को भी समान समस्याओं का सामना करना पड़ा था, जिसे फ्लैश-अटेंशन (FlashAttention) ने ऑनलाइन, फ्यूज्ड कर्नेल्स के माध्यम से हल किया था। हालांकि, अटेंशन डिस्टिलेशन पर इसे लागू करना गैर-तुच्छ (non-trivial) है क्योंकि KL डाइवर्जेंस एक वेटेड लोगिट-डिफरेंस (weighted logit-difference) टर्म के माध्यम से दो स्वतंत्र सॉफ्टमैक्स डिस्ट्रीब्यूशन को जोड़ता है, जिससे सिंगल-पास ऑनलाइन फॉर्मूलेशन गणितीय रूप से जटिल हो जाता है।

कार्यप्रणाली: StreamKL

StreamKL अटention KL डाइवर्जेंस के लिए पहला फ्यूज्ड, वन-पास GPU प्रिमिटिव पेश करता है जो मध्यवर्ती डिस्ट्रीब्यूशन के $O(N_Q N_K)$ मटेरियलाइजेशन को समाप्त करता है।

1. ऑनलाइन फॉर्मूलेशन (फॉरवर्ड पास)

StreamKL दो-डिस्ट्रीब्यूशन के कपल्ड KL रिडक्शन के लिए एक नया ऑनलाइन-अपडेट स्कीम व्युत्पन्न करता है। रो-वाइज (row-wise) KL डाइवर्जेंस को पुनर्गठित करके, लेखक इसे पांच प्रति-रो रनिंग स्कैलर्स के रूप में व्यक्त करते हैं:

$m_1, m_2$ : दो लोगिट डिस्ट्रीब्यूशन के लिए रनिंग मैक्सिमा।
$l_1, l_2$ : अननॉर्मलाइज्ड एक्सपोनेंशियल के योग (ऑनलाइन सॉफ्टमैक्स के समान)।
$acc$: एक एक््युमुलेटर (accumulator) जो वेटेड लोगिट डिफरेंस के लिए है।

पुनर्गठित ऑब्जेक्टिव है:
$L = \frac{acc}{l_1} + \text{LSE}_2 - \text{LSE}_1$
जहाँ $\text{LSE}_t = m_t + \log l_t$ है।

यह फॉर्मूलेशन कर्नेल को ऑन-चिप SRAM के माध्यम से क्वेरी-की टाइल्स (query-key tiles) को स्ट्रीम करने की अनुमति देता है, जिससे इन पांच स्कैलर्स को इंक्रीमेंटल रूप से अपडेट किया जा सकता है। महत्वपूर्ण रूप से, $acc $एक वेक्टर है जिसका आकार$ B_Q$ (ब्लॉक साइज) है, न कि एक पूर्ण मैट्रिक्स, जो स्टैंडर्ड फ्लैश-अटेंशन की तुलना में अरिथमेटिक रीस्केलिंग ओवरहेड को कम करता है। एल्गोरिदम कॉज़ल मास्किंग (causal masking) को हैंडल करने के लिए कॉज़ल फ्रंटियर से परे की की-टाइल्स (key tiles) को पूरी तरह से स्किप कर देता है।

2. रीकंप्यूटेशन स्ट्रैटेजी (बैकवर्ड पास)

बैकप्रोपैगेशन के दौरान क्वाड्रेटिक $P_1$ और $P_2$ मैट्रिसेस को स्टोर करने से बचने के लिए, StreamKL एक रीकंप्यूटेशन स्ट्रैटेजी का उपयोग करता है:

सेव की गई स्टेट (State Saved): केवल अंतिम KL वैल्यूज और लॉग-सम-एक्सप (LSE) वेक्टर्स ( $\text{LSE}_1, \text{LSE}_2$ ) को HBM ( $O(N_Q)$ ) में स्टोर किया जाता है।
रीजनरेशन (Regeneration): बैकवर्ड पास के दौरान, सेव किए गए LSEs और मूल $Q, K$ मैट्रिसेस का उपयोग करके अटेंशन प्रोबेबिलिटीज को टाइल-दर-टाइल पुन: उत्पन्न किया जाता है ( $P_i = e^{S_i - \text{LSE}}$ )।
ग्रेडिएंट कंप्यूटेशन: कर्नेल दो सेटिंग्स के लिए ग्रेडिएंट कंप्यूट करता है:
- सेटिंग 1: $P_1$ फिक्स्ड, $P_2$ को ऑप्टिमाइज़ करना (स्टैंडर्ड डिस्टिलेशन)।
- सेटिंग 2: $P_2$ फिक्स्ड, $P_1$ को ऑप्टिमाइज़ करना।
कर्नेल फ्यूजन (Kernel Fusion): लेखक एक फ्यूज्ड बैकवर्ड कर्नेल डिजाइन करते हैं जो प्रत्येक $QK$ टाइल पेयर को ठीक एक बार विज़िट करता है, जिससे HBM ट्रैफिक कम होता है। $dQ $(K के साथ) और$ dK$ (Q के साथ) को संचित करने के बीच संघर्ष को प्रबंधित करने के लिए, कर्नेल तब तक एटॉमिक ऑपरेशंस का उपयोग करता है जब तक कि बफर L2 कैश में फिट होता है, और जब कंटेंशन बहुत अधिक हो जाता है तो अलग-अलग कर्नेल्स पर स्विच कर जाता है।

3. GPU कर्नेल ऑप्टिमाइजेशन

स्प्लिट-K वेरिएंट (Split-K Variant): कम क्वेरी काउंट वाले वर्कलोड के लिए (जैसे ऑटो-रेग्रेसिव डिकोडिंग जहाँ $N_Q=1$ ), एक स्प्लिट-K वेरिएंट स्ट्रीमिंग मल्टीप्रोसेसरों (SMs) को सैचुरेट करने के लिए की-डायमेंशन को पार्टीशन करता है।
TMA एक्सेलरेशन (TMA Acceleration): हॉपर (Hopper) GPUs पर, कार्यान्वयन मेमोरी लेटेंसी को छिपाने के लिए एसिंक्रोनस बल्क मेमोरी कॉपी के लिए टेंसर मेमोरी एक्सेलेरेटर (TMA) का लाभ उठाता है।
न्यूमेरिकल स्टेबिलिटी (Numerical Stability): सेटिंग 2 में, न्यूमेरिकल इंस्टेबिलिटी ( $-\infty$ या NaN) से बचने के लिए लॉग-रेशियो कंप्यूटेशन को डायरेक्ट प्रोबेबिलिटी लॉग्स के बजाय लोगिट डिफरेंस और सेव किए गए LSEs का उपयोग करके किया जाता है।

मुख्य योगदान

बाधाओं की पहचान: पेपर उन क्वाड्रेटिक मेमोरी और IO बाधाओं की पहचान करता है जो लंबे कॉन्टेक्स्ट के लिए अटेंशन KL डाइवर्जेंस को स्केल करने से रोकते हैं।
ऑनलाइन फॉर्मूलेशन: यह अटेंशन KL डाइवर्जेंस के लिए पहले ऑनलाइन फॉर्मूलेशन को पेश करता है, जो एक पूरी तरह से फ्यूज्ड फॉरवर्ड प्रिमिटिव और एक रीकंप्यूटेशन-आधारित बैकवर्ड पास को सक्षम बनाता है जो क्वाड्रेटिक इंटरमीडिएट को समाप्त करता है।
कुशल कार्यान्वयन: लेखक दोनों फॉरवर्ड और बैकवर्ड पास के लिए समर्पित ऑप्टिमाइजेशन (Split-K, TMA, अडैप्टिव फ्यूजन) के साथ कुशल GPU कर्नेल्स डिजाइन और कार्यान्वित करते हैं।
स्केलेबिलिटी: StreamKL अटेंशन डिस्टिलेशन के अतिरिक्त HBM फुटप्रिंट को $O(N_Q N_K)$ से घटाकर $O(1)$ कर देता है, जिससे एकल GPU पर लंबे-कॉन्टेक्स्ट डिस्टिलेशन को सक्षम बनाया जा सका है।

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

NVIDIA A100 और H200 GPUs पर 4K से 512K तक के कॉन्टेक्स्ट लेंथ के लिए मूल्यांकन किया गया:

मेमोरी एफिशिएंसी (Memory Efficiency): StreamKL कॉन्टेक्स्ट लंबाई की परवाह किए बिना एक फ्लैट मेमोरी फुटप्रिंट बनाए रखता है। इसके विपरीत, बेसलाइन (PyTorch, torch.compile, FLA) क्वाड्रेटिक रूप से स्केल होते हैं और 16K–64K पर OOM एरर देते हैं। 64K पर, StreamKL torch.compile/FLA की तुलना में 16,000× से अधिक और PyTorch की तुलना में 135,000× से अधिक कम HBM का उपयोग करता है।
फॉरवर्ड लेटेंसी (Forward Latency):
- StreamKL नॉन-कॉज़ल सेटिंग्स में PyTorch की तुलना में 13–18× स्पीडअप और torch.compile/FLA की तुलना में 2.4–3.7× प्राप्त करता है।
- कॉज़ल मास्किंग के तहत, स्पीडअप बढ़कर PyTorch की तुलना में 26–43× और बेसलाइन्स की तुलना में 4.2–7.1× हो जाता है, क्योंकि StreamKL मास्क किए गए कंप्यूटेशन को पूरी तरह से स्किप करता है जबकि बेसलाइन्स पूर्ण मैट्रिसेस को मटेरियलाइज़ करते रहते हैं।
बैकवर्ड लेटेंसी (Backward Latency):
- सेटिंग 1 में, StreamKL (नॉन-कॉज़ल) की तुलना में PyTorch से 2.6–4.0× तेज़ है और (कॉज़ल) में 5.3–9.9× तेज़ है।
- सेटिंग 2 में, StreamKL (नॉन-कॉज़ल) की तुलना में PyTorch से 3.8–6.5× स्पीडअप और (कॉज़ल) में 9.1–14.0× स्पीडअप प्राप्त करता है।
- torch.compile की तुलना में, StreamKL प्रतिस्पर्धी या तेज़ है, विशेष रूप से कॉज़ल सेटिंग्स में जहाँ मास्क किए गए टाइल्स को स्किप करने से रीकंप्यूटेशन लागत आधी हो जाती है।
एब्लेशन (Ablation): स्प्लिट-K वेरिएंट छोटे $N_Q$ वर्कलोड के लिए 8.4× तक स्पीडअप प्रदान करता है। फ्यूज्ड बैकवर्ड कर्नेल छोटे $N_Q$ के लिए अलग-अलग कर्नेल्स की तुलना में 2.0× तक तेज़ है, जो अडैप्टिव सिलेक्शन स्ट्रैटेजी को सही ठहराता है।

महत्व

पेपर का दावा है कि StreamKL पहला समाधान है जो एक एकल GPU पर 64K और उससे अधिक कॉन्टेक्स्ट लेंथ के साथ अटेंशन डिस्टिलेशन वर्कलोड को बनाए रखने में सक्षम है। अटेंशन डिस्ट्रीब्यूशन के क्वाड्रेटिक मटेरियलाइजेशन को समाप्त करके, StreamKL उस प्राथमिक सिस्टम बाधा को हटा देता है जो आधुनिक LLM ट्रेनिंग (जैसे DeepSeek V3.2 और GLM-5 जैसे स्पार्स-अटेंशन मॉडल्स के लिए) में अटेंशन डिस्टिलेशन को स्केल करने से रोकती है। यह प्रदर्शित करता है कि फ्यूज्ड, ऑनलाइन प्रिमिटिव्स को KL डाइवर्जेंस जैसे जटिल कपल्ड डिस्ट्रीब्यूशन तक सफलतापूर्वक विस्तारित किया जा सकता है।