Apple Silicon के लिए ऑन-डिफ़ाइस अनुमान ऑप्टिमाइज़ करना
एक कस्टम स्थानीय इंजन जो प्रीफिल और डिकोड थ्रूपुट में सुधार करता है।
Apple silicon पर हाइब्रिड कंप्यूट क्लाउड में फ्रंटियर इंटेलिजेंस और Mac पर एक स्थानीय मॉडल के बीच एक कार्य का समन्वय करता है। क्लाउड मॉडल अनुसंधान और तर्क को संभालते हैं, जबकि एक स्थानीय मॉडल Mac पर निजी फाइलों और ऐप के साथ काम करता है।
श्रम के इस विभाजन को सहज महसूस करने के लिए, स्थानीय अनुमान को शेष कार्य के साथ गति बनाए रखनी चाहिए। इसके लिए एक ऐसे इंजन की आवश्यकता होती है जो प्रॉम्प्ट को तेज़ी से प्रोसेस कर सके और एक उच्च टोकन-generation दर बनाए रख सके।
Lily, हमारा हल्का स्थानीय अनुमान इंजन, विशेष रूप से Apple silicon और Qwen3.6-35B-A3B के लिए निर्मित है, जिसमें प्रीफिल और डिकोड के लिए अलग-अलग ऑप्टिमाइज़ेशन हैं। इंजन जल्द ही ओपन-सोर्स किया जाएगा।
परिचय
Mac पर LLMs चलाने का एक सामान्य तरीका MLX के साथ है, जो Apple silicon के लिए Apple का ओपन-सोर्स मशीन लर्निंग फ्रेमवर्क है। इसकी सहवर्ती लाइब्रेरी, MLX-LM, विभिन्न प्रकार के भाषा मॉडल के साथ टेक्स्ट को लोड करने और जनरेट करने के लिए आवश्यक घटक जोड़ती है। साथ में, MLX और MLX-LM स्थानीय LLM अनुमान के लिए एक ऑफ-द-शेल्फ, सामान्य-उद्देश्य स्टैक प्रदान करते हैं।
Qwen3.6-35B-A3B एक स्पार्स, हाइब्रिड मॉडल है: यह मिक्सचर-ऑफ-एक्सपर्ट्स (MoE) रूटिंग का उपयोग करता है और निश्चित आकार की रिकरेंट स्टेट्स को पूर्ण अटेंशन (full attention) के साथ जोड़ता है। ये आर्किटेक्चरल विकल्प आवश्यक गणना की मात्रा को कम करते हैं, लेकिन वे अनियमित वर्कलोड भी बनाते हैं। टोकन विभिन्न एक्सपर्ट वेट्स पर रूट होते हैं, और रिकरेंट स्टेट्स प्रकृति में क्रमिक होती हैं।
MLX-LM अनुमान चरणों और सामान्य वर्कलोड आकारों के लिए पहले से ही ऑप्टिमाइज़्ड कर्नेल का चयन करता है, लेकिन इसके पुनरुुपयोगी ऑपरेशनों को कई मॉडल आर्किटेक्चर का समर्थन करना चाहिए। Qwen को समर्पित एक इंजन मॉडल और रनटाइम स्तर पर विशेषज्ञता हासिल कर सकता है, जो मॉडल की निश्चित संरचना के इर्द-गिर्द कर्नेल, डेटा मूवमेंट और शेड्यूलिंग का समन्वय करता है।
Lily एक सिंगल प्रक्रिया में इस विशेषज्ञता को एंड-टू-एंड लागू करता है। एक Rust रनटाइम मॉडल चेकपॉइंट को लोड करता है और सत्र (session) स्टेट और जनरेशन लूप को प्रबंधित करता है, एक OpenAI-संगत चैट-कम्पलीशन API अनुरोधों और स्ट्रीम टोकन को स्वीकार करता है, और कस्टम मेटल कर्नेल Qwen-विशिष्ट ऑपरेशन निष्पादित करते हैं। निष्पादन पथ में न तो PyTorch है और न ही MLX है।

हम प्रीफिल और डिकोड प्रदर्शन को अलग-अलग मापते हैं। प्रीफिल थ्रूपुट कैप्चर करता है कि इंजन प्रॉम्प्ट को कितनी जल्दी प्रोसेस करता है; डिकोड थ्रूपुट कैप्चर करता है कि यह कितनी जल्दी आउटपुट टोकन जनरेट करता है।
हमने 40-कोर GPU और 128 GB यूनिफाइड मेमोरी वाले M5 Max द्वारा संचालित एक सिंगल MacBook Pro पर Qwen3.6-35B-A3B का बेंचिमार्च किया। प्रीफिल के लिए दस प्रॉम्प्ट लंबाई और डिकोड के लिए दस контекст लंबाई, 256 से 128K टोकन (K = 1,024) तक, इंजन औसतन 1.23× MLX-LM का प्रीफिल थ्रूपुट और 1.35× इसका डिकोड थ्रूपुट देता है। 4K-टोकन प्रॉम्प्ट और 4K-टोकन डिकोड контекст पर, कस्टम इंजन 5,749.9 प्रीफिल टोकन प्रति सेकंड और 186.6 डिकोड टोकन प्रति सेकंड तक पहुँचता है, जबकि MLX-LM के लिए यह 4,737.5 और 140.9 था। मल्टी-टर्न सत्र (multi-turn session) में, ये समय की बचत प्रत्येक अतिरिक्त मॉडल कॉल के साथ जुड़ती जाती है।

इसके बाद, हम बताते हैं कि कैसे Qwen का आर्किटेक्चर Apple silicon पर मॉडल-विशिष्ट ऑप्टिमाइज़ेशन के अवसर पैदा करता है। फिर हम परिणामी प्रीफिल और डिकोड परिवर्तनों से गुजरते हैं। MLX-LM के खिलाफ एंड-टू-एंड तुलना के साथ समाप्त करने से पहले हम यह भी कवर करते हैं कि अतिरिक्त ऑप्टिमाइज़ेशन कहाँ भुगतान करना बंद कर देता है।
Apple silicon पर Qwen-विशिष्ट ऑप्टिमाइज़ेशन के अवसर
Qwen तीन अलग-अलग वर्कलोड आकार बनाता है
Qwen3.6-35B-A3B में 35 अरब पैरामीटर हैं लेकिन प्रत्येक टोकन के लिए केवल लगभग 3 अरब सक्रिय होते हैं। एक राउटर 256 एक्सपर्ट सबनेटवर्क को स्कोर करता है और आठ का चयन करता है, साथ ही एक साझा एक्सपर्ट के साथ जो प्रत्येक टोकन को प्रोसेस करता है। यह स्पार्स MoE डिज़ाइन गणना को कम करता है लेकिन असमान कार्य उत्पन्न करता है: एक्सपर्ट्स को टोकन की अलग-अलग संख्या प्राप्त होती है, और प्रत्येक टोकन को एक्सपर्ट्स के एक अलग संयोजन से वेट्स की आवश्यकता होती है।
Qwen 10 फुल-अटेंशन लेयर को 30 Gated DeltaNet लेयर के साथ भी जोड़ता है। ये दोनों लेयर प्रकार अलग-अलग तरीकों से पहले की जानकारी को बनाए रखते हैं।
अटेंशन लेयर ग्रुपेड-क्वेरी अटेंशन (GQA) का उपयोग करती हैं। Qwen में 16 क्वेरी हेड और दो की-वैल्यू (KV) हेड हैं, जिसमें आठ क्वेरी हेड प्रत्येक KV हेड को साझा करते हैं। साझा करने से KV कैश छोटा हो जाता है और कैश्ड डेटा को क्वेरी हेड में पुनरुुपयोग करने की अनुमति मिलती है। कैश अभी भी प्रत्येक टोकन के लिए नई कुंजियों और मानों को संग्रहीत करता है, इसलिए जैसे-जैसे kontekst बढ़ता है प्रत्येक डिकोड चरण अधिक डेटा पढ़ता है।
इसके बजाय Gated DeltaNet पहले की जानकारी को एक निश्चित आकार की रिकरेंट स्टेट में संपीड़ित करता है। एक सीखा हुआ गेट नियंत्रित करता है कि मौजूदा स्टेट का कितना हिस्सा बनाए रखना है, जबकि एक डेल्टा अपडेट वर्तमान टोकन से जानकारी को शामिल करता है। मॉडल इन अपडेट को रिकरेंट रूप से परिभाषित करता है, इसलिए प्रत्येक टोकन पिछले टोकन द्वारा उत्पादित स्टेट पर निर्भर करता है। हालाँकि, प्रीफिल के दौरान, एक इंजन दो तरीकों से समान गणना का मूल्यांकन कर सकता है। यह स्टेट को आगे बढ़ाते हुए सीधे टोकन के माध्यम से स्कैन कर सकता है, या अपडेट को उन ब्लॉकों में पुनर्गठित कर सकता है जो अधिक मैट्रिक्स ऑपरेशन और टोकन-स्तरीय समानांतरता को उजागर करते हैं। कौन सा दृष्टिकोण तेज़ है यह मॉडल के आयामों, वर्कलोड और हार्डवेयर पर निर्भर करता है।
ये संरचनाएं मिलकर तीन कम्प्यूटेशनल पैटर्न बनाती हैं: असमान एक्सपर्ट समूह, बढ़ते हुए कैश पर अटेंशन, और एक निश्चित आकार की रिकरेंस जिसका मूल्यांकन सीधे या ब्लॉक में किया जा सकता है।
Apple silicon अलग-अलग वर्कलोड के लिए अलग-अलग पथ प्रदान करता है
प्रीफिल एक बार में कई प्रॉम्प्ट टोकन सक्रियण पंक्तियों को प्रोसेस करता है। यहाँ विचार किया गया स्थानीय वर्कलोड आम तौर पर एक समय में एक अनुरोध को डिकोड करता है (बैच 1) और प्रति चरण एक नई पंक्ति को प्रोसेस करता है। यह अंतर बदलता है कि समान मॉडल वेट्स का उपयोग कैसे किया जाता है। प्रीफिल सैकड़ों या हज़ारों पंक्तियों में वेट्स के प्रत्येक ब्लॉक का पुनरुुपयोग कर सकता है। डिकोड बड़े पैमाने पर नहीं कर सकता है, क्योंकि प्रत्येक नए टोकन के लिए वेट्स के माध्यम से एक और पास की आवश्यकता होती है।
Apple silicon CPU और GPU को यूनिफ़ाइड मेमोरी के पीछे रखता है, जो दोनों के लिए सुलभ एक सिंगल भौतिक मेमोरी पूल है। यह मॉडल को एक अलग GPU कॉपी को बनाए रखे बिना रेजिडेंट रहने की अनुमति देता है, लेकिन यह डेटा मूवमेंट को स्वतंत्र नहीं बनाता है। वेट्स और मध्यवर्ती मानों को पढ़ना अभी भी मेमोरी बैंडविड्थ का उपभोग करता है, जबकि रजिस्टर और अन्य ऑन-चिप स्टोरेज तेज़ हैं लेकिन बहुत छोटे हैं।
M5 GPU अलग-अलग कंप्यूट पथ भी प्रदान करता है। प्रीफिल की रैखिक लेयर जनरल मैट्रिक्स-मैट्रिक्स मल्टीप्लिकेशन (GEMM) का उपयोग करती हैं, जो एक बार में कई पंक्तियों पर एक वजन मैट्रिक्स लागू करती हैं। संगत GEMM Metal 4 टेंसर ऑपरेशनों के माध्यम से प्रत्येक GPU कोर में न्यूरल एक्सेलेरेटर का उपयोग कर सकते हैं। इसके बजाय Batch-1 डिकोड जनरल मैट्रिक्स-वेक्टर मल्टीप्लिकेशन (GEMV) का उपयोग करता है, जो एक पंक्ति पर समान वजन लागू करता है। कम वजन पुनरुुपयोग के साथ, GEMV मुख्य रूप से मेमोरी बैंडविड्थ द्वारा सीमित है और बड़े डेटा पुनरुुपयोग के साथ मैट्रिक्स संचालन के लिए डिज़ाइन किए गए न्यूरल एक्सेलेरेटर के बजाय GPU के वेक्टर अंकगणितीय तर्क इकाइयों (ALUs) के लिए अधिक उपयुक्त है।
ये निष्पादन पथ Lily के लिए अद्वितीय नहीं हैं। MLX समान यूनिफ़ाइड मेमोरी पर काम करता है और वर्कलोड के आकार के अनुसार ऑप्टिमाइज़्ड मैट्रिक्स और वेक्टर कर्नेल का चयन करता है। MLX-LM का Qwen कार्यान्वयन पहले से ही एक्सपर्ट कार्य को समूहीकृत करता है, फ्यूज्ड रिकरेंट मेटल कर्नेल के साथ Gated DeltaNet का मूल्यांकन करता है, और GQA-अवेयर अटेंशन का उपयोग करता है। ये क्षमताएं Apple silicon पर कुशल Qwen अनुमान के लिए साझा शुरुआती बिंदु हैं।
ऑप्टिमाइज़ेशन रणनीति
Lily का संकीर्ण दायरा इसे Qwen के सटीक आर्किटेक्चर और आयामों के इर्द-गिर्द इन साझा निष्पादन पथों (execution paths) को समन्वित करने की अनुमति देता है। यह चरण-विशिष्ट GPU पथों का उपयोग करता है, डेटा मूवमेंट को कम करने के लिए Qwen के एक्सपर्ट, रिकरेंट और अटेंशन वर्कलोड को मैप करता है, और मापे गए वर्कलोड के आकार से कर्नेल और लेआउट का चयन करता है। इस रणनीति के तीन भाग हैं:
- अनुमान चरण के लिए GPU पथ का मिलान करें। जब प्रीफिल कई पंक्तियों में वेट्स का पुनरुुपयोग कर सकता है तो मैट्रिक्स-उन्मुख निष्पादन का उपयोग करें, और जब batch-1 डिकोड एक समय में एक पंक्ति को प्रोसेस करता है तो वेक्टर-उन्मुख निष्पादन का उपयोग करें।
- डेटा मूवमेंट को कम करते हुए Qwen की संरचना को GPU पर मैप करें। उपयोग किए जाने तक वेट्स को संपीड़ित रखें, CPU पर वापस आए बिना रूट किए गए एक्सपर्ट कार्य को व्यवस्थित करें, अपने रिकरेंट स्कैन के माध्यम से Gated DeltaNet स्टेट को चिप पर बनाए रखें, और ग्रुपेड-क्वेरी अटेंशन द्वारा साझा किए गए KV डेटा का पुनरुुपयोग करें।
- वर्कलोड के आकार के अनुसार कर्नेल को अनुकूलित करें। प्रत्येक चरण के भीतर, उपलब्ध पंक्ति गणना, एक्सपर्ट्स में पंक्तियों के वितरण, ऑपरेशन के आयाम और वर्तमान kontekst लंबाई से टाइल के आकार, निष्पादन लेआउट और अटेंशन पथ का चयन करें।
निम्नलिखित अनुभाग इन विकल्पों की व्याख्या करते हैं। M5 Max पर मिलान किए गए एब्लेशन में मूल्यांकित ऑप्टिमाइज़ेशन के लिए, हम अन्यथा समान इंजन कॉन्फ़िगरेशन की तुलना करके उनके प्रभावों का अनुमान लगाते हैं जो केवल अध्ययन के तहत ऑप्टिमाइज़ेशन में भिन्न हैं। चूँकि ये प्रयोग हमारे इंजन के संस्करणों की तुलना स्वयं से करते हैं, इसलिए वे MLX-LM के खिलाफ अंतिम परिणामों को विघटित करने के बजाय तंत्र की व्याख्या करते हैं।
प्रीफिल: वेट्स का पुनरुुपयोग करें और GPU पर रूटिंग रखें
प्रीफिल एक बार में कई टोकन पंक्तियों को उजागर करता है, लेकिन Qwen उन पंक्तियों को एक्सपर्ट्स में असमान रूप से रूट करता है और अनुक्रम के माध्यम से रिकरेंट स्टेट को अपडेट करता है। इसके ऑप्टिमाइज़ेशन तीन समूहों में आते हैं: रूट की गई पंक्तियों के इर्द-गिर्द स्पार्स एक्सपर्ट कार्य को व्यवस्थित करें, Gated DeltaNet स्कैन को चिप पर रखें, और लंबे प्रॉम्प्ट को बाध्य (bounded) चंक में विभाजित करें।

स्पार्स एक्सपर्ट गणना को ऑप्टिमाइज़ करें
मैट्रिक्स गुणन के दौरान वेट्स को डीक्वांटाइज़ करें
Qwen3.6-35B-A3B चेकपॉइंट समूहवार एफ़िन 4-बिट परिमाणीकरण (क्वांटाइजेशन) का उपयोग करता है। प्रत्येक वजन को 4-बिट पूर्णांक कोड के रूप में संग्रहीत किया जाता है, जबकि 64 वजन के प्रत्येक समूह में bfloat16 स्केल और पूर्वाग्रह होता है जिसका उपयोग इसके मूल्यों को पुनर्गठित करने के लिए किया जाता है। यह 35-अरब-पैरामीटर मॉडल को लगभग 70 GB bfloat16 वेट्स से 19.4 GB चेकपॉइंट में कम कर देता है, जिससे मॉडल को Mac पर रेजिडेंट रखना व्यावहारिक हो जाता है।
मैट्रिक्स गुणन के लिए उपयोग किया जाने वाला Metal 4 टेंसर ऑपरेशन पैक किए गए 4-बिट प्रतिनिधित्व के बजाय bfloat16 ऑपरेंड का उपभोग करता है। गुणन से पहले, GPU को bfloat16 में वेट्स को पुनर्गठित करना होगा। Lily में ऑप्टिमाइज़्ड समूहीकृत GEMM इस रूपांतरण को एक समय में एक छोटे वेट टाइल पर करता है और रूट की गई सक्रियण पंक्तियों से इसे गुणा करने के लिए पर्याप्त लंबे समय तक ऑन-चिप थ्रेडग्रुप मेमोरी में परिणाम रखता है। संचयन 32-बिट फ़्लोटिंग पॉइंट का उपयोग करता है, और आउटपुट bfloat16 में लिखा जाता है। पूर्ण विस्तारित वेट सरणी यूनिफ़ाइड मेमोरी में कभी नहीं बनाई जाती है।
एब्लेशन में, डीक्वांटाइजेशन एक अलग ऑपरेशन के रूप में चलता है: यह यूनिफ़ाइड मेमोरी में 4-बिट वेट्स को bfloat16 सरणी में विस्तारित करता है, जिसके बाद मैट्रिक्स कर्नेल उस सरणी को वापस पढ़ता है। 512-टोकन प्रॉम्प्ट पर, समूहीकृत GEMM में डीक्वांटाइजेशन को ले जाने से इस मध्यवर्ती राइट and रीड को समाप्त करके एंड-टू-एंड प्रीफिल थ्रूपुट में 77.4% की वृद्धि हुई।
GPU पर एक्सपर्ट रूटिंग रखें
समूहीकृत GEMM को प्रत्येक एक्सपर्ट को असाइन की गई सक्रियण पंक्तियों को एक साथ संग्रहीत करने की आवश्यकता होती है। प्रति एक्सपर्ट आठ एक्सपर्ट का चयन करने के बाद, एक हिस्टोग्राम गिनती करता है कि प्रत्येक एक्सपर्ट को कितने असाइनमेंट मिले। एक प्रीफिक्स स्कैन उन गिनती को शुरुआती ऑफ़सेट में बदलता है, एक स्कैटर चरण पंक्तियों को उनके एक्सपर्ट समूहों में रखता है, और एक ब्लॉक मैप उन निश्चित आकार के मैट्रिक्स ब्लॉकों को सूचीबद्ध करता है जिन्हें समूहीकृत GEMM को प्रोसेस करना चाहिए।
ऑप्टिमाइज़्ड पथ प्रत्येक प्रॉम्प्ट चंक के लिए इस संपूर्ण अनुक्रम को एक कमांड बफर (GPU ऑपरेशनों का एक क्रमित बैच) में रखता है। इसके बजाय एक एब्लेशन रुकता है ताकि CPU रूटिंग मध्यवर्ती का निरीक्षण कर सके और अगले ऑपरेशन को सबमिट कर सके। हिस्टोग्राम और प्रीफिक्स स्कैन को GPU पर रखने से दो कर्नेल जुड़ जाते हैं लेकिन प्रत्येक MoE लेयर के अंदर CPU-GPU सिंक्रनाइज़ेशन हट जाता है।
512-टोकन प्रॉम्प्ट पर, GPU-रेसिडेंट रूटिंग को सक्षम करने से एंड-टू-एंड प्रीफिल में 89% की वृद्धि हुई। यह यह भी दिखाता है कि अकेला कर्नेल काउंट भ्रामक क्यों हो सकता है: तेज़ रास्ता अधिक कर्नेल लॉन्च करता है लेकिन लेयर के अंदर कभी भी CPU की प्रतीक्षा नहीं करता है।
टाइल के आकार को एक्सपर्ट लोड से मिलाएँ
2K-टोकन प्रॉम्प्ट पर, प्रत्येक टोकन को 256 एक्सपर्ट में से आठ पर रूट करने से 16,384 टोकन-एक्सपर्ट असाइनमेंट उत्पन्न होते हैं, या प्रति एक्सपर्ट औसतन 64 सक्रियण पंक्तियाँ होती हैं। वास्तविक वितरण असमान है: कुछ एक्सपर्ट्स को कई पंक्तियाँ मिलती हैं, जबकि अन्य को बहुत कम मिलती हैं।
समूहीकृत GEMM प्रत्येक एक्सपर्ट के आउटपुट को टाइलों में विभाजित करता है, जो मैट्रिक्स गुणन के आउटपुट के छोटे आयताकार ब्लॉक हैं। प्रत्येक टाइल को एक GPU थ्रेडग्रुप सौंपा गया है। Apple silicon के GPU पर, एक थ्रेडग्रुप में एक या अधिक सिमडग्रुप होते हैं, प्रत्येक में 32 थ्रेड होते हैं जो लॉकस्टेप में निर्देश निष्पादित करते हैं।
बड़े टाइल सेटअप लागत को अधिक पंक्तियों में फैलाते हैं और अधिक समानांतर कार्य को उजागर करते हैं, लेकिन जब किसी एक्सपर्ट को केवल कुछ पंक्तियाँ प्राप्त होती हैं तो बड़े टाइल का एक हिस्सा निष्क्रिय रहता है। इसलिए टाइल का आकार और सिमडग्रुप की संख्या जुड़े हुए हैं।
एक एब्लेशन टाइल को 16 पंक्तियों पर स्थिर करता है। उस नियंत्रण के खिलाफ, चार सिमडग्रुप के साथ 32-पंक्ति टाइल को सक्षम करने से 2K टोकन पर एंड-टू-एंड प्रीफिल में 13.2% सुधार हुआ।
चिप पर रिकरेंट स्टेट रखें
प्रीफिल के दौरान, प्रत्येक Gated DeltaNet लेयर अपनी रिकरेंट स्टेट को आगे बढ़ाते हुए क्रम में प्रॉम्प्ट को स्कैन करती है। रजिस्टर निवास अक्षम होने के साथ, एब्लेशन एक ब्लॉकवाइज स्कैन का उपयोग करता है। 2K-टोकन प्रॉम्प्ट पर, वह पथ प्रति लेयर 256 MiB (मेबिबाइट्स) स्टेट को मूव करता है और बार-बार बैरियर पर सहयोग करने वाले थ्रेड्स को रोकता है, ऐसे सिंक्रनाइज़ेशन बिंदु जहाँ सभी भाग लेने वाले थ्रेड्स को एक-दूसरे के लिए प्रतीक्षा करनी होगी।
रिकरेंट स्टेट एक मैट्रिक्स है। ऑप्टिमाइज़्ड कर्नेल प्रत्येक कॉलम को एक सिमडग्रुप (simdgroup) को निर्दिष्ट करता है। सिमडग्रुप कॉलम को अपने थ्रेड्स के बीच विभाजित करता है, कॉलम को उनके रजिस्टरों में एक बार लोड करता है, और पूरे स्कैन में स्टेट को आगे बढ़ाता है। थ्रेड्स थ्रेडग्रुप मेमोरी (थ्रेडग्रुप में साझा ऑन-चिप स्टोरेज) के बजाय सिमडग्रुप ऑपरेशन के माध्यम से मध्यवर्ती परिणामों का आदान-प्रदान करते हैं। स्कैन के बाद ही पूर्ण हुई स्टेट को वापस लिखा जाता है।
स्टेट और इसका गेट 32-बिट फ़्लोटिंग-पॉइंट फ़ॉर्मेट का उपयोग करते हैं क्योंकि क्रमिक अपडेट में छोटी राउंडिंग त्रुटियाँ बढ़ जाती हैं। क्वेरी और कुंजी सक्रियण bfloat16 में रहते हैं।
2K-टोकन प्रॉम्प्ट पर, रजिस्टर-रेसिडेंट स्कैन को सक्षम करने से एंड-टू-एंड प्रीफिल में 5.6% सुधार हुआ। एक्सपर्ट GEMM ने प्रीफिल समय का लगभग 90% हिस्सा लिया। क्रमिक स्कैन न्यूरल एक्सेलेरेटर से लाभ उठाने के लिए पर्याप्त पुनरुुपयोगी मैट्रिक्स कार्य को उजागर नहीं करता है।
प्रॉम्प्ट चंकिंग के साथ अस्थायी मेमोरी को बाध्य करें
रनटाइम एक लंबे प्रॉम्प्ट को एक साथ मेमोरी में प्रत्येक प्रॉम्प्ट टोकन के लिए अस्थायी डेटा रखने के बजाय बाध्य (bounded) चंक के अनुक्रम के रूप में प्रोसेस करता है। मॉडल वेट्स यूनिफ़ाइड मेमोरी में रेजिडेंट रहते हैं, जबकि रिकरेंट स्टेट और KV कैश एक चंक से दूसरे चंक तक контекस्ट ले जाते हैं। कोई भी पिछला контекस्ट त्याग नहीं जाता है।
चंकिंग के बिना, अस्थायी सक्रियण सरणी पूर्ण प्रॉम्प्ट के साथ बढ़ती है और यूनिफ़ाइड मेमोरी के लिए मॉडल वेट्स, रिकरेंट स्टेट और KV कैश के साथ प्रतिस्पर्धा करती है। चंकिंग एक समय में केवल एक सेगमेंट के अस्थायी मानों को लाइव रखती है, फिर अगले सेगमेंट को प्रोसेस करने से पहले उस स्टोरेज को रिलीज़ या पुनरुुपयोग करती है। यह पीक वर्किंग मेमोरी को सीमित करता है और इंजन को मॉडल के आउटपुट को बदले बिना लंबे प्रॉम्प्ट को प्रोसेस करने की अनुमति देता है।
चंक किया हुआ प्रीफिल कई इंजनों में लोकप्रिय है और इन मेमोरी-बाउंड वातावरणों में लंबे मल्टी-टर्न प्रक्षेपवक्र (trajectories) की सेवा के लिए महत्वपूर्ण है। अटेंशन लेयर के लिए कुल प्रीफिल समय पिछले चंक के बार-बार KV लोड से कुछ अतिरिक्त ओवरहेड के साथ, प्रॉम्प्ट लंबाई पर द्विघात (quadratic) रहता है।
डिकोड: प्रति टोकन मूव किए गए बाइट्स को कम करें
Batch-1 डिकोड एक समय में एक नई पंक्ति को प्रोसेस करता है। कम वजन पुनरुुपयोग के साथ, इसका थ्रूपुट मुख्य रूप से इस बात पर निर्भर करता है कि इंजन प्रत्येक टोकन के लिए कितने बाइट्स मूव करता है। डिकोड परिवर्तन चार समूहों में आते हैं: वन-रो वेट पाथ को ऑप्टिमाइज़ करें, प्रत्येक चरण को GPU पर रखें, मध्यवर्ती और स्टेट ट्रैफ़िक को कम करें, और अटेंशन कैश को कुशलता से पढ़ें।

वन-रो वेट पाथ को ऑप्टिमाइज़ करें
MLX पहले से ही विशेष मैट्रिक्स-वेक्टर कर्नेल के लिए वन-रो कार्य को प्रेषित करता है। चूंकि Lily MLX का उपयोग नहीं करता है, इसलिए कस्टम रनटाइम को समान बुनियादी रणनीति प्रदान करनी चाहिए। हमारा रो-पैरेलल GEMV एक सक्रियण पंक्ति के लिए डिज़ाइन किया गया है। वजन मैट्रिक्स के विभिन्न हिस्सों को समानांतर में पढ़ते समय एक सिमडग्रुप आउटपुट पर सहयोग करता है।
प्रत्येक डिकोड चरण को GPU पर रखें
GPU पर टोकन हैंडऑफ़ रखें
प्रत्येक डिकोड चरण अगले टोकन का चयन करके समाप्त होता है; निम्नलिखित चरण इनपुट के रूप में उस टोकन के साथ शुरू होता है। चयन को CPU और फिर वापस GPU पर भेजने से प्रत्येक टोकन में एक सिंक्रनाइज़ेशन बिंदु जुड़ जाता है। इसके बजाय हमारा रनटाइम दो कमांड बफर और दो GPU-रेसिडेंट टोकन स्लॉट के बीच वैकल्पिक होता है। GPU उच्चतम स्कोर वाले टोकन का चयन करता है और अगले डिकोड चरण के लिए सीधे इनपुट स्लॉट में अपनी टोकन ID लिखता है, जबकि CPU बाद के कार्य की तैयारी करता है।
स्वतंत्र GPU कार्य को ओवरलैप करें
एक दर्ज batch-1 डिकोड चरण में, एक टोकन जनरेट करने से 795 GPU कर्नेल लॉन्च हुए। उनकी निर्भरता ने 555 क्रमिक चरणों का गठन किया, जिससे कुछ कर्नेल समवर्ती रूप से चलाने के लिए स्वतंत्र हो गए। इसके बावजूद Metal के सीरियल निष्पादन मोड ने हर कर्नेल को क्रम में चलाया।
ऑप्टिमाइज़्ड डिकोड पथ एक समवर्ती Metal पास में वास्तविक डेटा निर्भरता को रिकॉर्ड करता है। जब GPU संसाधन अनुमति देते हैं तो स्वतंत्र कर्नेल लॉन्च एक ही समय में चल सकते हैं। एक बैरियर केवल तभी डाला जाता है जब बाद के कार्य को पिछले परिणाम की आवश्यकता होती है।
मध्यवर्ती और स्टेट ट्रैफ़िक को कम करें
अलग-अलग कर्नेल अक्सर एक मध्यवर्ती को भौतिक बनाते हैं: एक कर्नेल अस्थायी परिणाम को मेमोरी में लिखता है, और अगला परिणाम को वापस पढ़ता है। ऑप्टिमाइज़्ड डिकोड पथ चार श्रृंखलाओं को फ्यूज करता है: उनके गेटेड सक्रियण के साथ दो एक्सपर्ट इनपुट प्रोजेक्शन; इसके रूटिंग स्कोर और साझा-एक्सपर्ट परिणाम के साथ एक्सपर्ट आउटपुट प्रोजेक्शन; अटेंशन से पहले क्वेरी और कुंजी की तैयारी; और इसके सामान्यीकरण के साथ रिकरेंट अपडेट। प्रत्येक फ्यूज्ड कर्नेल अस्थायी मानों को मेमोरी के माध्यम से भेजने के बजाय रजिस्टरों में रखता है।
फ़्यूजन निर्भरता ग्राफ़ को भी छोटा करता है: जब एक मध्यवर्ती राइट गायब हो जाता है, तो वह बैरियर भी गायब हो जाता है जो इसके उपभोक्ता की रक्षा करता था।
अटेंशन कैश को कुशलता से पढ़ें
अटेंशन-कैश रीड्स को कोलेस करें
अटेंशन प्रत्येक डिकोड चरण के दौरान KV कैश से कुंजियाँ और मान पढ़ता है। एब्लेशन में, पड़ोसी GPU थ्रेड्स हमेशा पड़ोसी बाइट्स का अनुरोध नहीं करते हैं, जिससे मेमोरी सिस्टम को अधिक अलग-अलग लेनदेन की सेवा देने के लिए मजबूर होना पड़ता है। कोलेस किए गए लोड को सक्षम करने से आसन्न थ्रेड्स आसन्न बाइट्स का अनुरोध करते हैं ताकि हार्डवेयर अपने रीड्स को जोड़ सके।
bfloat16 कॉन्फ़िगरेशन पर, कोलेसिंग ने कुंजी बैंडविड्थ को 33.8 से 47.9 GB/s तक बढ़ा दिया, मान बैंडविड्थ को 42.0 से 61.8 GB/s तक बढ़ा दिया, और 3,840-टोकन контекст पर एंड-टू-एंड डिकोड में 2.1% सुधार किया।
KV पंक्तियों का पुनरुुपयोग करने के लिए क्वेरी हेड पैक करें
ग्रुपेड-क्वेरी अटेंशन आठ क्वेरी हेड को एक KV हेड साझा करने देता है। एब्लेशन में, प्रत्येक क्वेरी हेड एक अलग सिमडग्रुप (simdgroup) में चलता है, इसलिए सभी आठ स्वतंत्र रूप से एक ही कैश्ड KV रो का अनुरोध करते हैं। ऑप्टिमाइज़्ड कर्नेल चार क्वेरी हेड को एक थ्रेडग्रुप में पैक करता है, जो प्रत्येक KV रो को एक बार लोड करता है और चार अटेंशन गणनाओं में इसका पुनरुुपयोग करता है। दूसरा थ्रेडग्रुप शेष चार हेड को संभालता है।
यह तकनीक, जिसे आमतौर पर GQA पैकिंग कहा जाता है, समान अंकगणित करती है और आठ स्वतंत्र KV अनुरोधों को दो साझा लोड में कम करते हुए समान आउटपुट बाइट्स उत्पन्न करती है। unpacked एब्लेशन के खिलाफ, इसने 32K-टोकन контекस्ट पर एंड-टू-एंड डिकोड थ्रूपुट में 23.8% सुधार किया।
लंबे kontekst पर अटेंशन लेआउट स्विच करें
फुल-अटेंशन लेयर में प्रत्येक डिकोड चरण मौजूदा KV कैश को स्कैन करता है। एक फिक्स्ड-ब्लॉक लेआउट उस कैश को समान टुकड़ों में विभाजित करता है जिन्हें GPU समांतर रूप से प्रोसेस कर सकता है। जब कैश छोटा होता है तो इसका अतिरिक्त शेड्यूलिंग सार्थक नहीं होता है, लेकिन जैसे-जैसे контекст बढ़ता है, फिक्स्ड-ब्लॉक लेआउट काम को अधिक समान रूप से संतुलित करता है।
इस मॉडल के लिए, रनटाइम सामान्य अटेंशन पथ को 32K टोकन से नीचे रखता है और 32K या उससे अधिक पर फिक्स्ड-ब्लॉक पथ का उपयोग करता है। स्विच तब लागू होता है जब प्रत्येक हेड में 256 मान होते हैं और आठ क्वेरी हेड एक KV हेड साझा करते हैं; अन्य आकार सामान्य पथ पर रहते हैं। एक एब्लेशन इस स्विच को अक्षम करता है और हमेशा सामान्य पथ का उपयोग करता है। फिक्स्ड-ब्लॉक रूट को सक्षम करने से 32K पर एंड-टू-एंड डिकोड में 7.7%, 64K पर 27.4%, और 128K पर 40.2% सुधार हुआ।
अतिरिक्त ऑप्टिमाइज़ेशन की सीमाएं
कुछ परिवर्तनों ने एक अलग ऑपरेशन में सुधार किया लेकिन एंड-टू-एंड अनुमान में सुधार नहीं किया।
अनुमानित डिकोड (Speculative decoding), जो सत्यापित करने के लिए पूर्ण मॉडल के लिए टोकन का प्रस्ताव करने के लिए एक छोटे मॉडल का उपयोग करता है, ने batch-1 डिकोड को 18% धीमा कर दिया। सत्यापन ने दो से पांच पंक्तियों के समूहों को प्रोसेस किया, इस हार्डवेयर के लिए एक अकुशल आकार, और पंक्तियों ने अक्सर विभिन्न एक्सपर्ट्स का चयन किया, जिससे एक्सपर्ट-वेट डेटा पढ़ने की मात्रा बढ़ गई। ड्राफ्टर के आउटपुट शब्दावली को कम करने से ड्राफ्टर थ्रूपुट में 4.7–5.1% सुधार हुआ, लेकिन पूर्ण अनुमानित लूप तेज़ नहीं हुआ। यह परिणाम वर्कलोड-विशिष्ट है: हमारा बैच्ड Blackwell पर Qwen परिनियोजन विभिन्न स्थितियों में अनुमानित डिकोड का उपयोग करता है।
अन्य प्रयोगों में GPU लॉन्च को कम करना, संपूर्ण चरणों को ओवरलैप करना, बड़े प्रीफिल टाइलों का उपयोग करना, व्यापक फ़्यूजन लागू करना, राउटर को तेज़ी से चलाना, और टोकन चयन के साथ आउटपुट प्रोजेक्शन को जोड़ना शामिल था। किसी ने भी पूर्ण अनुमान लूप में सुधार नहीं किया।
हार्डवेयर सीमाओं के माप ने मुख्य प्रीफिल और डिकोड ऑपरेशनों में थोड़ा शेष हेडरूम भी दिखाया। MoE GEMM और GEMV अपने एक्सेस पैटर्न के लिए सबसे तेज़ निरंतर वेट-रीड दरों के 97.9% और 90.3% तक पहुँच गए। स्पार्स GEMV से अंकगणित को हटाने से थ्रूपुट केवल 0.2% बदल गया, जिससे यह पुष्टि हुई कि गणना के बजाय वेट रीड सीमित संसाधन थे। प्रीफिल का मैट्रिक्स गुणा इसी तरह अलगाव में सैद्धांतिक मैट्रिक्स सीमा के 93% और परीक्षण किए गए मॉडल के भीतर 80–86% तक पहुँच गया।
एंड-टू-एंड प्रदर्शन
एंड-टू-एंड तुलना दोनों इंजनों में समान 4-बिट चेकपॉइंट बाइट्स लोड करती है और एक 40-कोर, 128 GB M5 Max पर एक समय में एक अनुरोध चलाती है। प्रत्येक दौर के भीतर, पृष्ठभूमि भार और चिप तापमान में परिवर्तन से पूर्वाग्रह को कम करने के लिए दो इंजन वैकल्पिक क्रम में चलते हैं। हम MLX-LM के सबसे तेज़ डायरेक्ट-जेनरेशन पथ के खिलाफ तुलना करते हैं, इसके सर्वर के खिलाफ नहीं, इसलिए माप सर्विसिंग ओवरहेड के बजाय मॉडल निष्पादन पर केंद्रित है।
स्वीप में प्रीफिल के लिए दस प्रॉम्प्ट लंबाई और डिकोड के लिए दस kontekst लंबाई शामिल हैं, जो 256 से 128K टोकन तक हैं। प्रीफिल थ्रूपुट पहले बढ़ता है क्योंकि इंजन अधिक टोकन में निश्चित सेटअप लागत फैलाता है। प्रीफिल 4K-टोकन प्रॉम्प्ट के आसपास चरम पर होता है, फिर गिरता है क्योंकि जैसे-जैसे प्रॉम्प्ट बढ़ता है दस फुल-अटेंशन लेयर अधिक कार्य करती हैं। छोटे kontekst पर डिकोड लगभग सपाट रहता है और एक बार बढ़ते KV कैश को पढ़ना महत्वपूर्ण हो जाने पर घट जाता है। कस्टम इंजन हर दर्ज लंबाई पर तेज़ है।
चूंकि विशिष्ट निष्पादन फ़्लोटिंग-पॉइंट ऑपरेशन के क्रम को बदल सकता है, इसलिए हमने MLX-LM के खिलाफ संख्यात्मक स्थिरता की भी जाँच की। एक टीचर-फोर्स्ड तुलना में, दोनों इंजनों ने 192 स्थितियों में से प्रत्येक पर समान संदर्भ उपसर्ग (reference prefix) से अगले टोकन की भविष्यवाणी की, जिससे पहले के अंतरों को बाद के इनपुट को प्रभावित करने से रोका गया। Lily की पर्प्लेक्सिटी केवल 0.04% अधिक थी, और इसने परीक्षण की गई स्थितियों में से 96.35% पर समान शीर्ष-रैंक वाले टोकन का चयन किया।

स्थानीय प्लेटफ़ॉर्म के लिए निर्मित
Apple silicon एक छोटा डेटासेंटर GPU नहीं है। यह अपने स्वयं के हार्डवेयर और सॉफ़्टवेयर विशेषताओं के साथ एक संपूर्ण स्थानीय अनुमान प्लेटफ़ॉर्म है। यूनिफ़ाइड मेमोरी एक सिंगल नोड को इस बात पर बहुत उच्च सीमा देती है कि यह कितना मॉडल और स्टेट रख सकता है। M5 न्यूरल एक्सेलेरेटर प्रीफिल में सघन मैट्रिक्स कार्य को अवशोषित करते हैं। वेक्टर ALU डिकोड में बैंडविड्थ-बाउंड, कम-पुनरुुपयोग शेष को संभालते हैं।
Qwen विशेषज्ञता के लिए और अवसर जोड़ता है: GPU पर एक्सपर्ट रूटिंग और रिकरेंट स्टेट रखें, अनावश्यक मध्यवर्ती को समाप्त करें, स्वतंत्र कार्य को ओवरलैप करें, साझा KV डेटा का पुनरुुपयोग करें, और वर्कलोड के आकार के अनुसार कर्नेल को अनुकूलित करें।
मॉडल- और प्लेटफ़ॉर्म-विशिष्ट ऑप्टिमाइज़ेशन के साथ, एक Mac एक बड़े स्पार्स मॉडल को कुशलता से चला सकता है। भविष्य का काम मॉडल, चिप्स और सर्विसिंग वर्कलोड में कवरेज का विस्तार करेगा, और एक कॉन्फ़िगरेशन पर यहाँ मान्य किए गए तंत्र को अधिक सामान्य रनटाइम नीति में बदल देगा।
व्यापक सिद्धांत इंजन को मॉडल के आर्किटेक्चर और हार्डवेयर के विशिष्ट कंप्यूट और मेमोरी पथ दोनों से मिलाना है। जैसे-जैसे फ्रंटियर ओपन-वेट मॉडल और हार्डवेयर विकसित होते हैं, उच्च-प्रदर्शन स्थानीय अनुमान तेजी से उन इंजनों पर निर्भर करेगा जो उनके मतभेदों को दूर करने वाले इंजनों के बजाय दोनों के अनुकूल हैं।