🤖 AI

FlashMLA-ETAP: Efficient Transpose Attention Pipeline for Accelerating MLA Inference on NVIDIA H20 GPUs

यह शोध पत्र FlashMLA-ETAP प्रस्तुत करता है, जो एक कुशल ट्रांसपोज़ अटेंशन पाइपलाइन (Efficient Transpose Attention Pipeline) का उपयोग करने वाला एक नवीन ढांचा है, जो WGMMA ऑपरेशन्स को अनुकूलित करके NVIDIA H20 GPU पर मल्टी-हेड लेटेंट अटेंशन (Multi-Head Latent Attention) इन्फरेंस को महत्वपूर्ण रूप से तेज करता है, जिससे उच्च संख्यात्मक स्थिरता बनाए रखते हुए मौजूदा विधियों की तुलना में पर्याप्त गति प्राप्त होती है।

मूल लेखक: Pengcuo Dege, Qiuming Luo, Rui Mao, Chang Kong

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

📖 4 मिनट में पढ़ें☕ कॉफ़ी ब्रेक में पढ़ें

CC BY 4.0

मूल लेखक: Pengcuo Dege, Qiuming Luo, Rui Mao, Chang Kong

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

कल्पना कीजिए कि आप एक बहुत लंबी किताब (संदर्भ/context) पढ़ने की कोशिश कर रहे हैं ताकि एक छोटे से सवाल (प्रश्न/query) का उत्तर दिया जा सके। यह मूल रूप से वही है जो DeepSeek-R1 जैसा एक बड़ा AI मॉडल टेक्स्ट जेनरेट करते समय करता है: वह अगला शब्द क्या कहना है, यह तय करने के लिए अब तक पढ़ी गई हर चीज़ को देखता है।

यह शोध पत्र एक नई विधि पेश करता है जिसे FlashMLA-ETAP कहा जाता है, जो AI को उस किताब के पन्नों को पलटने का एक स्मार्ट तरीका सिखाने जैसा है, विशेष रूप से NVIDIA H20 जैसे कंप्यूटर चिप के लिए।

यहाँ समस्या और समाधान का रोज़मर्रा के उदाहरणों का उपयोग करके विवरण दिया गया है:

समस्या: "अजीब शेल्फ" वाली समस्या

NVIDIA H20 GPU को एक लाइब्रेरी में काम करने वाले लाइब्रेरियन के रूप में सोचें, जिसके पास एक विशिष्ट नियम है: शेल्फ को स्थिर रहने के लिए कम से कम 64 किताबों चौड़ा होना चाहिए। यदि आप शेल्फ पर केवल 16 किताबें रखने की कोशिश करते हैं, तो लाइब्रेरियन को खाली जगह को नकली किताबों (पैडिंग) से भरना पड़ता है ताकि शेल्फ गिरे नहीं। इससे समय और ऊर्जा बर्बाद होती है।

AI की दुनिया में, "किताबें" अटेंशन हेड्स (मस्तिष्क के वे हिस्से जो अलग-अलग चीजों पर ध्यान केंद्रित करते हैं) हैं। जब एक ही सर्वर पर 8 H20 GPU के साथ विशाल DeepSeek-R1 मॉडल को चलाया जाता है, तो काम को विभाजित किया जाता है। प्रत्येक GPU अंततः केवल 16 हेड्स को संभालता है।

चूंकि 16, आवश्यक 64 से कम है, इसलिए H20 GPU को अपने हार्डवेयर नियमों को संतुष्ट करने के लिए बहुत सारा "नकली काम" (पैडिंग) करना पड़ता है। यह ऐसा है जैसे लाइब्रेरियन को 16 असली बक्सों को ले जाने के लिए 64 खाली बक्से ले जाने के लिए मजबूर किया जा रहा हो। यह AI को धीमा और अक्षम बनाता है, खासकर जब वह "किताब" जिसे वह पढ़ रहा है बहुत लंबी (लॉन्ग कॉन्टेक्स्ट) हो लेकिन वह जिस सवाल का जवाब दे रहा है वह बहुत छोटा (एक समय में एक शब्द) हो।

समाधान: किताब को तिरछा करना (ETAP)

लेखकों ने एक तकनीक बनाई है जिसे ETAP (एफिशिएंट ट्रांसपोज़ अटेंशन पाइपलाइन) कहा जाता है। 16 हेड्स को 64-चौड़ी शेल्फ के नियम में फिट करने की कोशिश करने के बजाय, वे समस्या को तिरछा (sideways) कर देते हैं।

कल्पित करें कि एक छोटी पंक्ति में 16 हेड्स को देखने के बजाय, आप किताब की लंबाई (जो हजारों पन्ने हो सकती है) को मुख्य आयाम (dimension) के रूप में देखते हैं। चूंकि किताब बहुत लंबी है, इसलिए यह बिना किसी नकली किताब के "64-चौड़ी शेल्फ" की आवश्यकता को आसानी से भर देती है।

आयामों को बदलकर—बातचीत के लंबे इतिहास को मुख्य "चौड़ाई" और छोटे सवाल को "ऊंचाई" मानकर—H20 GPU बिना किसी खाली पैडिंग के बर्बाद किए पूर्ण गति से काम कर सकता है। यह ऐसा है जैसे यह महसूस करना कि 16 छोटी वस्तुओं को एक बड़े बॉक्स में फिट करने के बजाय, आपको उन्हें लंबवत (vertically) स्टैक करना चाहिए जहाँ पर्याप्त जगह है।

परिणाम: तेज़ और अधिक सटीक

शोध पत्र का दावा है कि यह "तिरछा" दृष्टिकोण H20 GPU पर एक बड़ा अंतर पैदा करता है:

बहुत तेज़: लंबी बातचीत की लंबाई (64,000 शब्द) पर, FlashMLA-ETAP इस विशिष्ट मॉडल (FlashMLA) के लिए पिछले सबसे अच्छे तरीके की तुलना में 2.78 गुना तेज़ है। यह FlashAttention-3 और FlashInfer जैसे अन्य लोकप्रिय तरीकों की तुलना में भी काफी तेज़ है।
अधिक सटीक: आमतौर पर, जब आप AI गणनाओं को तेज़ करने की कोशिश करते हैं, तो आप थोड़ी सटीकता खो सकते हैं (जैसे संख्याओं को बहुत आक्रामक रूप से राउंड करना)। हालाँकि, यह नई विधि वास्तव में FlashAttention-3 की तुलना में अधिक सटीक है, जिसमें त्रुटि दर (RMSE) बहुत कम है। यह किताब को तेज़ी से पढ़ने और उसे बेहतर ढंग से समझने जैसा है।

यह क्यों महत्वपूर्ण है

अधिकांश AI अनुकूलन (optimizations) बहुत महंगे, हाई-एंड चिप्स (जैसे H100) के लिए डिज़ाइन किए गए हैं। H20 एक "मिड-टियर" चिप है—अच्छी है, लेकिन सबसे शक्तिशाली नहीं है। यह शोध पत्र दिखाता है कि सही सॉफ्टवेयर ट्रिक (ETAP) के साथ, आप मिड-टियर चिप्स को विशिष्ट कार्यों के लिए बहुत बेहतर प्रदर्शन करने के लिए बना सकते हैं। यह एक मानक सेडान को एक विशिष्ट प्रकार की सड़क पर स्पोर्ट्स कार की तरह कुशलता से चलाने का तरीका खोजने जैसा है, बस गियर बदलने के तरीके को बदलकर।

संक्षेप में: FlashMLA-ETAP एक सॉफ्टवेयर अपडेट है जो NVIDIA H20 GPU को यह बताता है कि लंबी AI बातचीत को कैसे पढ़ना है, जिससे बेकार की मेहनत खत्म होती है और AI तेज़ी से और अधिक सटीकता से जवाब देता है।

तकनीकी सारांश: FlashMLA-ETAP

समस्या विवरण
NVIDIA H20 जैसे मिड-टियर GPUs पर DeepSeek-R1 (671 बिलियन पैरामीटर्स) जैसे बड़े-स्तरीय मॉडल्स की तैनाती, मल्टी-हेड लेटेंट अटेंशन (MLA) इन्फरेंस के लिए महत्वपूर्ण दक्षता संबंधी चुनौतियां पेश करती है। जबकि हाई-एंड GPUs (जैसे H100, H800) में मौजूदा अनुकूलन को संभालने के लिए पर्याप्त कंप्यूट क्षमता होती है, NVIDIA H20 की मध्यम 148 TFLOPS FP16 कंप्यूट क्षमता एक बाधा उत्पन्न करती है। विशेष रूप से, सिंगल-इंस्टेंस डिप्लॉयमेंट परिदृश्यों में जहाँ मॉडल के 128 अटेंशन हेड्स को 8 GPUs में विभाजित किया जाता है (परिणामस्वरूप प्रति GPU 16 हेड्स), प्रति-GPU हेड काउंट WarpGroup Matrix-Multiply-Accumulate (WGMMA) निर्देशों की न्यूनतम M-डायमेंशन आवश्यकता (64) से कम हो जाता है। यह बेमेल (mismatch) अनावश्यक पैडिंग (padding) को आवश्यक बनाता है, जिससे उन डिकोडिंग चरणों के दौरान कंप्यूट उपयोगिता 25% से नीचे गिर जाती है जो छोटे क्वेरी लेंथ और लंबे की-वैल्यू (KV) कॉन्टेक्स्ट लेंथ द्वारा विशेषता वाले होते हैं। मौजूदा फ्रेमवर्क्स जैसे FlashAttention-3 और FlashMLA, जो हाई-एंड हार्डवेयर के लिए अनुकूलित हैं, इन आर्किटेक्चरल बाधाओं और सब-ऑप्टिमल मेमोरी एक्सेस पैटर्न के कारण H20 के संसाधनों का पूर्ण उपयोग करने में विफल रहते हैं।

कार्यप्रणाली (Methodology)
इन अक्षमताओं को दूर करने के लिए, लेखक FlashMLA-ETAP प्रस्तावित करते हैं, जो FlashMLA आर्किटेक्चर में एफिशिएंट ट्रांसपोज़ अटेंशन पाइपलाइन (ETAP) को एकीकृत करता है।

एफिशिएंट ट्रांसपोज़ अटेंशन पाइपलाइन (ETAP): ETAP ट्रांसपोज़ेशन का लाभ उठाकर कम्प्यूटेशनल डायमेंशन को H20 के WGMMA बाधाओं के साथ संरेखित करके मानक अटेंशन गणना को पुनर्गठित करता है। पारंपरिक MLA इन्फरेंस में, गणना $S = Q \cdot K^T$ , $P = \text{softmax}(S)$ , और $O = P \cdot V$ का पालन करती है। ETAP इस प्रक्रिया को इस प्रकार पुनरorient करता है:
1. $S^T = K \cdot Q^T \in \mathbb{R}^{N \times N_q}$
2. $P^T = \text{softmax}(S^T) \in \mathbb{R}^{N \times N_q}$
3. $O = V^T \cdot P^T \in \mathbb{R}^{d \times N_q}$
4. $O = O^T$
WGMMA ऑपरेशन्स में लंबे KV कॉन्टेक्स्ट लेंथ ( $N$ ) को M-डायमेंशन और छोटे क्वेरी लेंथ ( $N_q$ ) को N-डायमेंशन मानकर, ETAP क्वेरी डायमेंशन पर अनावश्यक पैडिंग की आवश्यकता को समाप्त करता है। यह संरेखण ऑटोरेग्रेसिव डिकोडिंग में क्वेरी और KV लेंथ के बीच विषमता का लाभ उठाता है, जिससे कम्प्यूटेशनल ओवरहेड काफी कम हो जाता है।
एकीकरण और कार्यान्वयन: FlashMLA-ETAP, FlashMLA के कोऑपरेटिव थ्रेड एरे (CTA) फ्रेमवर्क के भीतर ETAP को लागू करता है। यह सर्कुलर शेयर्ड मेमोरी बफर के साथ प्रोड्यूसर-कन्ज्यूमर सिंक्रोनाइज़ेशन और ब्लॉक-वाइज प्रोसेसिंग का उपयोग करता है। एल्गोरिदम गणना को दो वारपग्रुप्स के बीच विभाजित करता है: एक कंज्यूमर वारपग्रुप जो अटेंशन स्कोर और सॉफ्टमैक्स सांख्यिकी (statistics) को संभालता है, और एक प्रोड्यूसर वारपग्रुप जो हाई-बैंडविड्थ मेमोरी (HBM) से शेयर्ड मेमोरी (SMEM) में डेटा लोडिंग का प्रबंधन करता है। यह डिज़ाइन KV कैश न्यूनीकरण के लिए FlashMLA के लो-रैंक जॉइंट कंप्रेशन को बनाए रखते हुए, H20 के लिए मैट्रिक्स मल्टीप्लिकेशन और सॉफ्टमैक्स चरणों को अनुकूलित करता है।
सैद्धांतिक सामान्यीकरण (Theoretical Generalizability): पेपर एक सैद्धांतिक विश्लेषण प्रदान करता है जो दर्शाता है कि ETAP की ट्रांसपोज़िशन रणनीति FlashAttention-3 और FlashInfer सहित अन्य अटेंशन फ्रेमवर्क्स के साथ संगत है। चूंकि ये फ्रेमवर्क्स मौलिक अटेंशन कम्प्यूटेशन पैटर्न साझा करते हैं, इसलिए ETAP को उनके टिलिंग और कर्नेल फ्यूजन रणनीतियों में न्यूनतम संशोधनों के साथ एकीकृत किया जा सकता है, जो संभावित रूप से अन्य मिड-टियर GPU आर्किटेक्चर तक इसके लाभों का विस्तार करता है।

प्रमुख योगदान

ETAP का परिचय: एक नवीन कम्प्यूटेशन मोड जो मिड-टियर GPUs पर WGMMA पैडिंग आवश्यकताओं के कारण होने वाली अनावश्यक गणनाओं को कम करने के लिए ट्रांसपोज़िशन का उपयोग करता है।
FlashMLA-ETAP का विकास: FlashMLA का एक उन्नत संस्करण जो H20 पर कुशल MLA इन्फरेंस सक्षम करने के लिए ETAP को शामिल करता है, जिसका डिज़ाइन FlashAttention-3 और FlashInfer जैसे अन्य फ्रेमवर्क्स में सहज एकीकरण के लिए बनाया गया है।
प्रदर्शन सत्यापन: H20 GPU पर बेसलाइन फ्रेमवर्क्स की तुलना में महत्वपूर्ण स्पीडअप और बेहतर संख्यात्मक स्थिरता (numerical stability) प्रदर्शित करने वाले प्रायोगिक परिणाम।

परिणाम
प्रयोग NVIDIA H20 GPU पर DeepSeek-R1 मॉडल का उपयोग करके किए गए थे, जिसमें 512 से 64K तक के सीक्वेंस लेंथ और 16 और 32 के बैच साइज के लिए FlashMLA-ETAP की तुलना FlashMLA, FlashAttention-3 और FlashInfer से की गई।

स्पीडअप: 64K के सीक्वेंस लेंथ और 16 के बैच साइज पर, FlashMLA-ETAP ने FlashMLA (89 TFLOPS/s बनाम 32 TFLOPS/s) पर 2.78x स्पीडअप प्राप्त किया। इसने FlashAttention-3 से 5.24x और FlashInfer से 4.94x अधिक प्रदर्शन भी किया। बैच साइज 32 के लिए, इसने 64K पर FlashMLA पर 2.72x स्पीडअप प्राप्त किया।
स्केलेबिलिटी: ETAP का प्रदर्शन लाभ सीक्वेंस लेंथ के साथ बढ़ता गया, जो लॉन्ग-कॉन्टेक्स्ट परिदृश्यों में पैडिंग ओवरहेड को कम करने में इसकी प्रभावशीलता को उजागर करता है।
संख्यात्मक स्थिरता: FlashMLA-ETAP ने FlashAttention-3 की तुलना में बेहतर संख्यात्मक सटीकता प्रदर्शित की। डबल-प्रिसिजन रेफरेंस के विरुद्ध रूट मीन स्क्वायर एरर (RMSE) FlashMLA-ETAP के लिए $1.25 \times 10^{-5}$ था, जबकि FlashAttention-3 के लिए यह $1.9 \times 10^{-4}$ था, जो त्रुटि में 15.2x की कमी को दर्शाता है।

महत्व
पेपर का दावा है कि FlashMLA-ETAP मिड-टियर GPUs के लिए एक स्केलेबल समाधान प्रदान करके संसाधन-बाधित इन्फरेंस में एक महत्वपूर्ण अंतर को संबोधित करता है। NVIDIA H20 के विशिष्ट आर्किटेक्चरल बाधाओं के लिए अनुकूलित करके, यह फ्रेमवर्क सिंगल-इंस्टेंस परिदृश्यों में DeepSeek-R1 जैसे बड़े मॉडल्स की कुशल तैनाती को सक्षम बनाता है। लेखकों का मानना है कि ETAP का सामान्यीकरण योग्य डिज़ाइन हार्डवेयर-अवेयर ऑप्टिमाइज़ेशन के व्यापक उपयोग का मार्ग प्रशस्त करता है, जो संभावित रूप से विविध GPU आर्किटेक्चर के लिए कुशल इन्फरेंस को लोकतांत्रिक बनाने और कम्प्यूटेशनल दक्षता के लिए भविष्य के अनुसंधान का समर्थन करने में मदद करेगा।

समस्या: "अजीब शेल्फ" वाली समस्या

समाधान: किताब को तिरछा करना (ETAP)

परिणाम: तेज़ और अधिक सटीक

यह क्यों महत्वपूर्ण है

इसी तरह के और पेपर