Apple सिलिकॉनसाठी ऑन-डिव्हाइस इन्फरन्स ऑप्टिमाइझ करणे

Apple सिलिकॉनवर हायब्रिड संगणन क्लाउडमधील सीमा बुद्धिमत्ता आणि मॅकवरील स्थानिक मॉडेल यांच्यातील कार्याचे समन्वय साधते. क्लाउड मॉडेल्स संशोधन आणि तर्क हाताळतात, तर स्थानिक मॉडेल मॅकवरील खाजगी फाइल्स आणि अ‍ॅप्ससह कार्य करते.

कामाच्या या विभागाला अखंड वाटण्यासाठी, स्थानिक इन्फरन्सने उर्वरित कार्याशी सुसंगत असले पाहिजे. त्यासाठी अशा इंजिनची आवश्यकता आहे जे प्रॉम्प्ट्सवर जलद प्रक्रिया करू शकेल आणि उच्च टोकन-जनरेशन दर टिकवून ठेवू शकेल.

लिली, आमचे हलके वजनी स्थानिक इन्फरन्स इंजिन, विशेषतः Apple सिलिकॉन आणि Qwen3.6-35B-A3B साठी तयार केले गेले आहे, ज्यामध्ये प्रिफिल आणि डिकोडसाठी स्वतंत्र ऑप्टिमायझेशन आहेत. हे इंजिन लवकरच ओपन-सोर्स केले जाईल.

परिचय

मॅकवर LLM चालवण्याचा एक सामान्य मार्ग म्हणजे MLX सह, जे Apple सिलिकॉनसाठी Apple चे ओपन-सोर्स मशीन लर्निंग फ्रेमवर्क आहे. त्याचे सहकारी लायब्ररी, MLX-LM, विस्तृत श्रेणीतील भाषा मॉडेल्ससह मजकूर लोड करण्यासाठी आणि जनरेट करण्यासाठी आवश्यक घटक जोडते. एकत्रितपणे, MLX आणि MLX-LM स्थानिक LLM इन्फरन्ससाठी ऑफ-द-शेल्फ, सामान्य-हेतू स्टॅक प्रदान करतात.

Qwen3.6-35B-A3B हे एक विरल, संकरित मॉडेल आहे: हे मिक्सचर-ऑफ-एक्सपर्ट्स (MoE) रुटिंग वापरते आणि पूर्ण लक्ष देऊन निश्चित-आकाराची पुनरावृत्ती होणारी राज्ये एकत्र करते. हे वास्तुशिल्पीय पर्याय आवश्यक संगणनाची मात्रा कमी करतात, परंतु ते अनियमित वर्कलोड देखील तयार करतात. टोकन्स वेगवेगळ्या तज्ज्ञ वजनांवर रूट होतात, आणि पुनरावृत्ती होणारी राज्ये स्वभावाने अनुक्रमिक असतात.

MLX-LM इन्फरन्स टप्प्यांसाठी आणि सामान्य वर्कलोड आकारांसाठी आधीपासूनच ऑप्टिमाइझ केलेले कर्नल निवडतो, परंतु त्याच्या पुनर्वापर करण्यायोग्य ऑपरेशन्सना अनेक मॉडेल आर्किटेक्चरचे समर्थन करावे लागते. Qwen ला समर्पित इंजिन मॉडेल आणि रनटाइम स्तरावर विशेषीकरण करू शकते, मॉडेलच्या निश्चित रचनेभोवती कर्नल, डेटा हालचाल आणि शेड्यूलिंगचे समन्वय साधू शकते.

लिली एकाच प्रक्रियेमध्ये हे विशेषीकरण अंत ते अंत लागू करते. Rust रनटाइम मॉडेल चेकपॉइंट लोड करतो आणि सत्र स्थिती आणि जनरेशन लूप व्यवस्थापित करतो, 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 आहे. मल्टी-टर्न सत्रात, प्रत्येक अतिरिक्त मॉडेल कॉलसह ही वेळ बचत जमा होते.

पुढे, Qwen चे वास्तुकला Apple सिलिकॉनवर मॉडेल-विशिष्ट ऑप्टिमायझेशन संधी कशी तयार करते ते आम्ही स्पष्ट करतो. त्यानंतर आम्ही परिणामी प्रिफिल आणि डिकोड बदलांमधून मार्गक्रमण करतो. MLX-LM विरुद्ध एंड-टू-एंड तुलनेने बंद करण्यापूर्वी अतिरिक्त ऑप्टिमायझेशन कुठे मोबदला देणे थांबवते हे देखील आम्ही कव्हर करतो.

Apple सिलिकॉनवर Qwen-विशिष्ट ऑप्टिमायझेशनच्या संधी

Qwen तीन भिन्न वर्कलोड आकार तयार करतो

Qwen3.6-35B-A3B मध्ये 35 अब्ज पॅरामीटर्स आहेत परंतु प्रत्येक टोकनसाठी केवळ सुमारे 3 अब्ज सक्रिय होतात. एक राउटर 256 तज्ज्ञ सबनेटवर्क्सचे गुण मोजतो आणि आठ निवडतो, प्रत्येक टोकनवर प्रक्रिया करणारा एक सामायिक तज्ज्ञ सोबत. ही विरल MoE रचना संगणन कमी करते परंतु असमान कार्य तयार करते: तज्ज्ञांना टोकन्सची वेगवेगळी संख्या प्राप्त होते, आणि प्रत्येक टोकनला तज्ज्ञांच्या वेगवेगळ्या संयोजनातून वजनांची आवश्यकता असते.

Qwen हे 10 फुल-अटेंशन लेअर्स 30 गेेटेड डेल्टानेट लेअर्ससह देखील एकत्र करते. हे दोन लेयर प्रकार वेगवेगळ्या मार्गांनी आधीची माहिती कायम ठेवतात.

अटेंशन लेअर्स ग्रुपड्-क्वे अटेंशन (GQA) वापरतात. Qwen मध्ये 16 क्वे हेड आणि दोन की-व्हॅल्यू (KV) हेड आहेत, प्रत्येक KV हेड सामायिक करणारे आठ क्वे हेड आहेत. सामायिक केल्याने KV कॅशे लहान होतो आणि कॅश केलेला डेटा क्वे हेडवर पुन्हा वापरला जाण्याची परवानगी मिळते. कॅशे अजूनही प्रत्येक टोकनसाठी नवीन की आणि मूल्ये संग्रहित करतो, त्यामुळे संदर्भ वाढल्याने प्रत्येक डिकोड स्टेप अधिक डेटा वाचतो.

गेेटेड डेल्टानेट त्याऐवजी आधीची माहिती निश्चित-आकाराच्या पुनरावृत्ती स्थितीत संकुचित करतो. विद्यमान स्थितीपैकी किती ठेवायची हे एक शिकलेले गेर नियंत्रित करते, तर डेल्टा अद्यतन वर्तमान टोकनवरील माहिती समाविष्ट करते. मॉडेल ही अद्यतने पुनरावृत्तीने परिभाषित करते, त्यामुळे प्रत्येक टोकन पूर्ववर्ती टोकन द्वारे उत्पादित स्थितीवर अवलंबून असते. प्रिफिल दरम्यान, तथापि, एक इंजिन समान संगणनाचे दोन मार्गांनी मूल्यमापन करू शकते. हे स्थिती पुढे नेत असताना टोकन्सद्वारे थेट स्कॅन करू शकते, किंवा अधिक मॅट्रिक्स ऑपरेशन्स आणि टोकन-स्तरीय समांतरता उघड करणाऱ्या ब्लॉकमध्ये अद्यतनांची पुनर्रचना करू शकते. कोणता दृष्टिकोन वेगवान आहे हे मॉडेल परिमाण, वर्कलोड आणि हार्डवेअरवर अवलंबून असते.

एकत्र येऊन, ही रचना तीन संगणकीय नमुने तयार करतात: असमान तज्ज्ञ गट, वाढत्या कॅशेवर लक्ष केंद्रित करणे, आणि एक निश्चित-आकाराची पुनरावृत्ती जी थेट किंवा ब्लॉकमध्ये मूल्यमापन केली जाऊ शकते.

Apple सिलिकॉन वेगवेगळ्या वर्कलोड्ससाठी वेगवेगळे मार्ग प्रदान करतो

प्रिफिल एकाच वेळी अनेक प्रॉम्प्ट टोकन सक्रियता ओळींवर प्रक्रिया करते. येथे विचार केला गेलेला स्थानिक वर्कलोड सामान्यत: एका वेळी एक विनंती डिकोड करतो (बच 1) आणि प्रति स्टेप एक नवीन ओळ प्रक्रिया करतो. हा फरक समान मॉडेल वजन कसे वापरले जाते हे बदलतो. प्रिफिल शेकडो किंवा हजारो ओळींवर वजनांच्या प्रत्येक ब्लॉकम्याचा पुनर्वापर करू शकतो. डिकोड मोठ्या प्रमाणावर करू शकत नाही, कारण प्रत्येक नवीन टोकनसाठी वजनांमधून आणखी एक पास आवश्यक असतो.

Apple सिलिकॉन CPU आणि GPU ला युनिफाइड मेमरीच्या मागे ठेवते, दोन्हीसाठी प्रवेशयोग्य एकच भौतिक मेमरी पूल. यामुळे वेगळी GPU प्रत कायम न ठेवता मॉडेल निवासी राहू शकते, परंतु यामुळे डेटा हालचाल मोफत होत नाही. वजन आणि मध्यवर्ती मूल्ये वाचणे अजूनही मेमरी बँडविड्थ वापरते, तर रजिस्टर आणि इतर ऑन-चिप स्टोरेज वेगवान परंतु खूप लहान आहेत.

M5 GPU देखील वेगवेगळे संगणन मार्ग प्रदान करते. प्रिफिलचे रेषीय लेयर जनरल मॅट्रिक्स-मॅट्रिक्स मल्टिप्लिकेशन (GEMM) वापरतात, एकाच वेळी अनेक ओळींवर वजन मॅट्रिक्स लागू करतात. सुसंगत GEMM Metal 4 टेंसर ऑपरेशन्स द्वारे प्रत्येक GPU कोरमधील न्यूरल अ‍ॅक्सिलेटर वापरू शकतात. बॅच-1 डिकोड त्याऐवजी जनरल मॅट्रिक्स-वेक्टर मल्टिप्लिकेशन (GEMV) वापरते, एका ओळीवर समान वजन लागू करते. कमी वजन पुनवापरासह, GEMV प्रामुख्याने मेमरी बँडविड्थद्वारे मर्यादित असते आणि अधिक डेटा पुनवापरासह मॅट्रिक्स ऑपरेशन्ससाठी डिझाइन केलेल्या न्यूरल अ‍ॅक्सिलेटरपेक्षा GPU च्या वेक्टर अंकगणित तर्क युनिट्स (ALUs) साठी अधिक योग्य आहे.

हे अंमलबजावणी मार्ग केवळ लिलीपुरता मर्यादित नाहीत. MLX समान युनिफाइड मेमरीवर कार्य करते आणि वर्कलोड आकारानुसार ऑप्टिमाइझ केलेले मॅट्रिक्स आणि वेक्टर कर्नल निवडते. MLX-LM चे Qwen अंमलबजावणी आधीपासूनच तज्ज्ञ कार्य गटबद्ध करते, फ्यूज्ड पुनरावृत्ती मेटल कर्नलसह गेेटेड डेल्टानेटचे मूल्यमापन करते, आणि GQA-अवेर अटेंशन वापरते. हे क्षमता Apple सिलिकॉनवर कार्यक्षम Qwen इन्फरन्ससाठी सामायिक सुरूवातीचा बिंदू आहेत.

ऑप्टिमायझेशन धोरण

लिलीचा अरुंद वा ̧वर या सामायिक अंमलबजावणी मार्गांना Qwen च्या अचूक वास्तुकला आणि परिमाणांभोवती समन्वय साधण्यास अनुमती देतो. हे टप्पा-विशिष्ट GPU मार्ग वापरते, डेटा हालचाल कमी करण्यासाठी Qwen चे तज्ज्ञ, पुनरावृत्ती आणि लक्ष वर्कलोड्स मॅप करते, आणि मोजलेल्या वर्कलोड आकारावरून कर्नल आणि लेआउट निवडते. धोरणाचे तीन भाग आहेत:

1. इन्फरन्स टप्प्याशी GPU मार्ग जुळवा. जेव्हा प्रिफिल अनेक ओळींवर वजन पुन्हा वापरू शकते तेव्हा मॅट्रिक्स-ориентированная अंमलबजावणी वापरा, आणि जेव्हा बॅच-1 डिकोड एका वेळी एक ओळ प्रक्रिया करते तेव्हा वेक्टर-ориентированная अंमलबजावणी वापरा.
2. डेटा हालचाल कमी करताना Qwen ची रचना GPU वर मॅप करा. वापरेपर्यंत वजन संकुचित ठेवा, CPU वर परत न जाता रुट केलेले तज्ज्ञ कार्य आयोजित करा, गेेटेड डेल्टानेट स्थिती त्याच्या पुनरावृत्ती स्कॅनद्वारे चिपवर ठेवा, आणि ग्रुपड्-क्वे अटेंशन द्वारे सामायिक केलेला KV डेटा पुन्हा वापरा.
3. वर्कलोड आकारानुसार कर्नल अनुकूल करा. प्रत्येक टप्प्यात, उपलब्ध ओळ गणनेतून, तज्ज्ञांमधील ओळींचे वितरण, ऑपरेशनचे परिमाण आणि वर्तमान संदर्भ लांबीमधून टाइल आकार, अंमलबजावणी लेआउट आणि अटेंशन मार्ग निवडा.

खालील विभाग हे पर्याय स्पष्ट करतात. M5 Max वर जुळलेल्या आब्लेशनमध्ये मूल्यांकन केलेल्या ऑप्टिमायझेशनसाठी, आम्ही अभ्यासात असलेल्या ऑप्टिमायझेशनमध्येच भिन्न असलेल्या अन्यथा समान इंजिन कॉन्फिगरेशनची तुलना करून त्यांच्या प्रभावाचा अंदाज लावतो. हे प्रयोग आमच्या इंजिनच्या आवृत्त्यांची स्वतःशी तुलना करत असल्यामुळे, ते MLX-LM विरुद्ध अंतिम परिणामांचे विघटन करण्याऐवजी यंत्रणा स्पष्ट करतात.

प्रिफिल: वजन पुन्हा वापर करा आणि रुटिंग GPU वर ठेवा

प्रिफिल एकाच वेळी अनेक टोकन ओळी उघड करतो, परंतु Qwen त्या ओळी तज्ज्ञांवर असमान रीतीने रुट करतो आणि अनुक्रमाद्वारे पुनरावृत्ती स्थिती अद्यतनित करतो. त्याचे ऑप्टिमायझेशन तीन गटांमध्ये येतात: रुट केलेल्या ओळींभोवती विरल तज्ज्ञ कार्य आयोजित करा, गेेटेड डेल्टानेट स्कॅन चिपवर ठेवा, आणि लांब प्रॉम्प्ट्सना मर्यादित चंक्समध्ये विभाजित करा.

विरल तज्ज्ञ संगणन ऑप्टिमाइझ करा

मॅट्रिक्स गुणाकारादरम्यान वजन डीक्वांटाईज करा

Qwen3.6-35B-A3B चेकपॉइंट समूहवाईज आफाइन 4-बिट परिमाण वापरतो. प्रत्येक वजन 4-बिट पूर्णांक कोड म्हणून संग्रहित केले जाते, तर 64 वजनांच्या प्रत्येक समूहाला त्याची मूल्ये पुनर्रचित करण्यासाठी वापरले जाणारे bfloat16 स्केल आणि पूर्वाग्रह सामायिक करतो. यामुळे 35-अब्ज-पॅरामीटर मॉडेल अंदाजे 70 GB bfloat16 वजनावरून 19.4 GB चेकपॉइंटमध्ये कमी होते, ज्यामुळे मॉडेल मॅकवर निवासी ठेवणे व्यावहारिक बनते.

मॅट्रिक्स गुणाकारासाठी वापरले जाणारे मेटल 4 टेंसर ऑपरेशन पॅक केलेल्या 4-बिट प्रतिनिधित्वासह bfloat16 ऑपरेंड्स वापरते. गुणाकारापूर्वी, GPU ने bfloat16 मध्ये वजन पुनर्रचित करणे आवश्यक आहे. लिली मधील ऑप्टिमाइझ केलेले समूह GEMM हे रूपांतरण एका वेळी एक लहान वजन टाइल करते आणि रूट केलेल्या सक्रियता ओळींनी गुणाकार करण्यासाठी ते ऑन-चिप थ्रेडग्रुप मेमरीमध्ये तेवढा वेळ ठेवते जितका वेळ आवश्यक असतो. एकत्रीकरण 32-बिट फ्लोटिंग पॉइंट वापरते, आणि आउटपुट bfloat16 मध्ये लिहिले जाते. संपूर्ण विस्तारित वजन सरणी युनिफाइड मेमरीमध्ये कधीही तयार केली जात नाही.

आब्लेशनमध्ये, डीक्वांटायझेशन स्वतंत्र ऑपरेशन म्हणून चालते: हे 4-बिट वजनांना युनिफाइड मेमरीमधील bfloat16 सरणीमध्ये विस्तृत करते, त्यानंतर मॅट्रिक्स कर्नल ती सरणी परत वाचतो. 512-टोकन प्रॉम्प्टवर, डीक्वांटायझेशन समूह GEMM मध्ये हलवल्याने या मध्यवर्ती लेख आणि वाचणे दूर करून एंड-टू-एंड प्रिफिल थ्रूपुट 77.4% ने वाढला.

तज्ज्ञ रुटिंग GPU वर ठेवा

समूह GEMM ला प्रत्येक तज्ज्ञाला नियुक्त केलेल्या सक्रियता ओळी एकत्र संग्रहित करण्याची आवश्यकता असते. प्रति तज्ज्ञ आठ तज्ज्ञ निवडल्यानंतर, हिस्टोग्राम मोजतो की प्रत्येक तज्ज्ञाला किती असाइनमेंट गेले. एक प्रिफिक्स स्कॅन त्या मोजणीला सुरुवातीच्या ऑफसेटमध्ये बदलते, एक स्कॅटर स्टेप ओळींना त्यांच्या तज्ज्ञ गटांमध्ये ठेवते, आणि एक ब्लॉक मॅप निश्चित-आकाराचे मॅट्रिक्स ब्लॉक सूचीबद्ध करतो ज्यावर समूह GEMM ने प्रक्रिया करणे आवश्यक आहे.

ऑप्टिमाइझ केलेला मार्ग प्रत्येक प्रॉम्प्ट चंकसाठी ही संपूर्ण अनुक्रम एकाच कमांड बफरमध्ये, GPU ऑपरेशन्सच्या क्रमाने जोडलेल्या बॅचमध्ये ठेवतो. त्याऐवजी एक आब्लेशन थांबते जेणेकरून CPU रुटिंग इंटरमीडिएट्सची तपासणी करू शकेल आणि पुढील ऑपरेशन सबमिट करू शकेल. GPU वर हिस्टोग्राम आणि प्रिफिक्स स्कॅन ठेवल्याने दोन कर्नल जोडले जातात परंतु प्रत्येक MoE लेयरमधील CPU-GPU सिंक्रोनायझेशन काढून टाकले जाते.

512-टोकन प्रॉम्प्टवर, GPU-रेसिडेंट रुटिंग सक्षम केल्याने एंड-टू-end प्रिफिलमध्ये 89% वाढ झाली. हे हे देखील दर्शवते की केवळ कर्नल संख्या दिशाभूल करणारी असू शकते: वेगवान मार्ग अधिक कर्नल लॉन्च करतो परंतु लेयरमधील CPU ची कधीही वाट पाहत नाही.

तज्ज्ञ भाराशी टाइल आकार जुळवा

2K-टोकन प्रॉम्प्टवर, प्रत्येक टोकनला 256 पैकी आठ तज्ज्ञांकडे रुट केल्याने 16,384 टोकन-तज्ज्ञ असाइनमेंट मिळतात, किंवा प्रति तज्ज्ञ सरासरी 64 सक्रियता ओळी मिळतात. प्रत्यक्ष वितरण असमान आहे: काही तज्ज्ञांना अनेक ओळी प्राप्त होतात, तर इतरांना कमी मिळतात.

समूह GEMM प्रत्येक तज्ज्ञाचे आउटपुट टाइल्समध्ये विभाजित करते, जे मॅट्रिक्स गुणाकाराच्या आउटपुटचे लहान आयताकृती ब्लॉक आहेत. प्रत्येक टाइल एका GPU थ्रेडग्रुपला नियुक्त केली जाते. Apple सिलिकॉनच्या GPU वर, थ्रेडग्रुपमध्ये एक किंवा अधिक सिमद्ग्रुप असतात, प्रत्येक 32 थ्रेड्सचा समावेश असतो जे लॉकस्टेपमध्ये सूचना कार्यान्वित करतात.

मोठ्या टाइल्स सेटअप खर्च अधिक ओळींवर पसरवतात आणि अधिक समांतर कार्य उघड करतात, परंतु जेव्हा एखाद्या तज्ज्ञाला फक्त काही ओळी प्राप्त होतात तेव्हा मोठ्या टाइलचा एक भाग निष्क्रिय राहतो. त्यामुळे टाइलचा आकार आणि सिमद्ग्रुपची संख्या जोडलेली असते.

एक आब्लेशन टाइल 16 ओळींवर निश्चित करतो. त्या नियंत्रणाविरुद्ध, चार सिमद्ग्रुपसह 32-ओळ टाइल सक्षम केल्याने 2K टोकन्सवर एंड-टू-एंड प्रिफिलमध्ये 13.2% सुधारणा झाली.

पुनरावृत्ती स्थिती चिपवर ठेवा

प्रिफिल दरम्यान, प्रत्येक गेेटेड डेल्टानेट लेयर त्याची पुनरावृत्ती स्थिती पुढे नेत असताना प्रॉम्प्ट क्रमाने स्कॅन करतो. रजिस्टर रेजिडेन्सी अक्षम करून, आब्लेशन ब्लॉकवाइज स्कॅन वापरते. 2K-टोकन प्रॉम्प्टवर, तो मार्ग प्रति लेयर 256 MiB (मेबिबाइट्स) स्थिती हलवतो आणि वारंवार अडथळ्यांवर, सिंक्रोनायझेशन पॉईंट्सवर जिथे सर्व सहभागी थ्रेड्सनी एकमेकांची वाट पाहिली पाहिजे तिथे सहकार्य करणारे थ्रेड्स थांबवतो.

पुनरावृत्ती राज्य एक मॅट्रिक्स आहे. ऑप्टिमाइझ केलेले कर्नल प्रत्येक कॉलमला एका सिमद्ग्रुपला नियुक्त करते. सिमद्ग्रुप कॉलमला त्याच्या थ्रेड्समध्ये विभाजित करतो, कॉलमला त्यांच्या रजिस्टरमध्ये एकदा लोड करतो आणि संपूर्ण स्कॅनमध्ये स्थिती घेऊन जातो. थ्रेड्स थ्रेडग्रुप मेमरीऐवजी, थ्रेडग्रुपवर सामायिक केलेल्या ऑन-चिप स्टोरेजऐवजी, सिमद्ग्रुप ऑपरेशन्सद्वारे मध्यवर्ती परिणाम एक्सचेंज करतात. पूर्ण झालेली स्थिती स्कॅन नंतरच परत लिहिली जाते.

स्थिती आणि त्याची गेर 32-बिट फ्लोटिंग-पॉइंट फॉरमॅट वापरतात कारण अनुक्रमिक अद्यतनांमध्ये लहान गोलाकार त्रुटी एकत्र होतात. क्वे आणि की सक्रियता bfloat16 मध्ये राहतात.

2K-टोकन प्रॉम्प्टवर, रजिस्टर-रेसिडेंट स्कॅन सक्षम केल्याने एंड-टू-एंड प्रिफिलमध्ये 5.6% सुधारणा झाली. तज्ज्ञ GEMMs ने प्रिफिल वेळेच्या सुमारे 90% भाग घेतला. अनुक्रमिक स्कॅन न्यूरल अ‍ॅक्सिलेटरचा फायदा घेण्यासाठी पुरेसे पुनर्वापर करण्यायोग्य मॅट्रिक्स कार्य उघड करत नाही.

प्रॉम्प्ट चंकिंगसह तात्पुरती मेमरी मर्यादित करा

रनटाइम प्रत्येक प्रॉम्प्ट टोकनसाठी तात्पुरता डेटा मेमरीमध्ये एकाच वेळी ठेवण्याऐवजी लांब प्रॉम्प्टला मर्यादित चंक्सच्या अनुक्रम म्हणून प्रक्रिया करतो. मॉडेल वजन युनिफाइड मेमरीमध्ये राहतात, तर पुनरावृत्ती स्थिती आणि KV कॅशे एका चंकवरून पुढील चंकपर्यंत संदर्भ वाहून नेतात. कोणताही आधीचा संदर्भ टाकून दिला जात नाही.

चंकिंगशिवाय, तात्पुरती सक्रियता सरणी संपूर्ण प्रॉम्प्टसह वाढतात आणि युनिफाइड मेमरीसाठी मॉडेल वजन, पुनरावृत्ती स्थिती आणि KV कॅशे यांच्याशी स्पर्धा करतात. चंकिंग एका वेळी फक्त एकाच विभागाची तात्पुरती मूल्ये थेट ठेवते, नंतर पुढील विभागावर प्रक्रिया करण्यापूर्वी तो साठा रिलीज करतो किंवा पुन्हा वापरतो. हे पीक वर्किंग मेमरी कॅप करते आणि मॉडेलचे आउटपुट न बदलता इंजिनला लांब प्रॉम्प्ट्सवर प्रक्रिया करण्याची परवानगी देते.

चंक्ड प्रिफिल अनेक इंजिनमध्ये लोकप्रिय आहे आणि या मेमरी-मर्यादित वातावरणात लांब मल्टी-टर्न ट्रॅजेक्टरीज सर्व्ह करण्यासाठी महत्त्वपूर्ण आहे. आधीच्या चंक्सच्या वारंवार KV लोडमुळे काही जोडलेल्या ओव्हरहेडसह, अटेंशन लेअर्ससाठी एकूण प्रिफिल वेळ प्रॉम्प्ट लांबीवर चतुर्भुज राहतो.

डिकोड: प्रति टोकन हलवलेले बाइट्स कमी करा

बॅच-1 डिकोड एका वेळी एक नवीन ओळ प्रक्रिया करतो. कमी वजन पुनवापरासह, त्याचा थ्रूपुट प्रामुख्याने इंजिन प्रत्येक टोकनसाठी किती बाइट्स हलवते यावर अवलंबून असतो. डिकोड बदल चार गटांमध्ये येतात: एक-ओळ वजन मार्ग ऑप्टिमाइझ करा, प्रत्येक स्टेप GPU वर ठेवा, मध्यवर्ती आणि राज्य रहदारी कमी करा, आणि अटेंशन कॅशे कार्यक्षमतेने वाचा.

एक-ओळ वजन मार्ग ऑप्टिमाइझ करा

MLX आधीच विशेष मॅट्रिक्स-वेक्टर कर्नलना एक-ओळ कार्य पाठवते. लिली MLX वापरत नसल्यामुळे, सानुकूल रनटाइमने समान मूलभूत धोरण प्रदान करणे आवश्यक आहे. आमचा ओळ-समांतर GEMV एका सक्रियता ओळीसाठी डिझाइन केला आहे. वजन मॅट्रिक्सचे विविध भाग समांतरपणे वाचत असताना एक सिमद्ग्रुप आउटपुटवर सहकार्य करतो.

प्रत्येक डिकोड स्टेप GPU वर ठेवा

टोकन हँडऑफ GPU वर ठेवा

प्रत्येक डिकोड स्टेप पुढील टोकन निवडून समाप्त होते; पुढील स्टेप त्या टोकनसह इनपुट म्हणून सुरू होते. CPU ला आणि नंतर GPU ला निवड पाठवल्याने प्रत्येक टोकनमध्ये सिंक्रोनायझेशन पॉइंट जोडला जातो. आमचा रनटाइम त्याऐवजी दोन कमांड बफर आणि दोन GPU-रेसिडेंट टोकन स्लॉट्समध्ये पर्यायी ठरतो. GPU सर्वोच्च-स्कोअरिंग टोकन निवडतो आणि पुढील डिकोड स्टेपसाठी इनपुट स्लॉटमध्ये त्याचा टोकन आयडी थेट लिहितो, तर CPU त्यानंतरचे कार्य तयार करतो.

स्वतंत्र GPU कार्य आच्छादित करा

एका रेकॉर्ड केलेल्या बॅच-1 डिकोड स्टेपमध्ये, टोकन जनरेट केल्याने 795 GPU कर्नल लॉन्च झाले. त्यांच्या परावलंबनांनी 555 अनुक्रमिक टप्पे तयार केले, ज्यामुळे काही कर्नल समवर्तीपणे चालण्यासाठी मोकळे झाले. तरीही मेटलच्या अनुक्रमिक अंमलबजावणी मोडने प्रत्येक कर्नल क्रमाने चालवला.

ऑप्टिमाइझ केलेला डिकोड मार्ग एका समवर्ती मेटल पासमध्ये वास्तविक डेटा परावलंबन रेकॉर्ड करतो. जेव्हा GPU संसाधने परवानगी देतात तेव्हा स्वतंत्र कर्नल लॉन्च एकाच वेळी चालू शकतात. जेव्हा नंतरच्या कार्यासाठी आधीच्या निकालाची आवश्यकता असते तेव्हाच अडथळा घातला जातो.

मध्यवर्ती आणि राज्य रहदारी कमी करा

वेगवेगळे कर्नल अनेकदा मध्यवर्ती साहित्य भौतिक करतात: एक कर्नल तात्पुरता परिणाम मेमरीमध्ये लिहितो, आणि पुढील परिणाम परत वाचतो. ऑप्टिमाइझ केलेला डिकोड मार्ग चार साखळ्या फ्यूज करतो: गेेटेड सक्रियतेसह दोन तज्ज्ञ इनपुट प्रोजेक्शन; त्याच्या रुटिंग स्कोअरसह तज्ज्ञ आउटपुट प्रोजेक्शन आणि सामायिक-तज्ज्ञ परिणाम; अटेंशनपूर्वी क्वे आणि की तयारी; आणि त्याच्या सामान्यीकरणासह पुनरावृत्ती अद्यतन. प्रत्येक फ्यूज्ड कर्नल तात्पुरती मूल्ये मेमरीद्वारे पाठवण्याऐवजी रजिस्टरमध्ये ठेवतो.

फ्युजन मुळे डिपेंडन्सी ग्राफ देखील लहान होतो: जेव्हा एखादा मध्यवर्ती लेख नष्ट होतो, तेव्हा त्याच्या ग्राहकाचे संरक्षण करणारा अडथळा देखील नष्ट होतो.

अटेंशन कॅशे कार्यक्षमतेने वाचा

अटेंशन-कॅशे वाचन कोलेस करा

प्रत्येक डिकोड स्टेप दरम्यान अटेंशन KV कॅशेमधून की आणि व्हॅल्यू वाचतो. आब्लेशनमध्ये, शेजारील GPU थ्रेड्स नेहमीच शेजारील बाइट्सची विनंती करत नाहीत, ज्यामुळे मेमरी सिस्टमला अधिक स्वतंत्र व्यवहारांची सेवा देणे भाग पडते. कोलेस केलेले लोड सक्षम केल्याने लगतचे थ्रेड लगतच्या बाइट्सची विनंती करतात जेणेकरून हार्डवेअर त्यांच्या वाचण्यांची एकत्र करू शकेल.

bfloat16 कॉन्फिगरेशनवर, कोलेसिंगने की बँडविड्थ 33.8 वरून 47.9 GB/s पर्यंत वाढवली, व्हॅल्यू बँडविड्थ 42.0 वरून 61.8 GB/s पर्यंत वाढवली, आणि 3,840-टोकन संदर्भावर एंड-टू-एंड डिकोडमध्ये 2.1% सुधारणा केली.

KV ओळी पुन्हा वापरण्यासाठी क्वे हेड पॅक करा

ग्रुपड्-क्वे अटेंशन आठ क्वे हेडना एक KV हेड सामायिक करू देते. आब्लेशनमध्ये, प्रत्येक क्वे हेड वेगळ्या सिमद्ग्रुपमध्ये चालते, त्यामुळे सर्व आठ स्वतंत्रपणे समान कॅश केलेल्या KV ओळीची विनंती करतात. ऑप्टिमाइझ केलेले कर्नल चार क्वे हेडना एका थ्रेडग्रुपमध्ये पॅक करते, जे प्रत्येक KV ओळी एकदा लोड करते आणि चार अटेंशन गणनेवर ती पुन्हा वापरते. दुसरे थ्रेडग्रुप उर्वरित चार हेड हाताळते.

हे तंत्र, ज्याला सामान्यतः GQA पॅकिंग म्हटले जाते, आठ स्वतंत्र KV विनंती दोन सामायिक लोडमध्ये कमी करताना समान अंकगणित करते आणि समान आउटपुट बाइट्स तयार करते. अनपॅक केलेल्या आब्लेशनविरुद्ध, त्याने 32K-टोकन संदर्भावर एंड-टू-एंड डिकोड थ्रूपुट 23.8% ने सुधारले.

लांब संदर्भांवर अटेंशन लेआउट स्विच करा

पूर्ण-अटेंशन थरांमधील प्रत्येक डिकोड पायरी विद्यमान KV कॅशे स्कॅन करते. एक निश्चित-ब्लॉक लेआउट त्या कॅशेला समान तुकड्यांमध्ये विभाजित करतो ज्या GPU समांतरपणे प्रक्रिया करू शकते. जेव्हा कॅशे लहान असतो तेव्हा त्याचे अतिरिक्त नियोजन फायदेशीर नसते, परंतु जसजاس संदर्भ वाढतो तसतसे निश्चित-ब्लॉक लेआउट कार्य अधिक समान रीतीने संतुलित करतो.

या मॉडेलसाठी, रनटाइम सामान्य अटेंशन मार्ग 32K टोकन्सच्या खाली ठेवतो आणि 32K किंवा त्याहून अधिक वर निश्चित-ब्लॉक मार्ग वापरतो. जेव्हा प्रत्येक हेडमध्ये 256 मूल्ये असतात आणि आठ क्वे हेड KV हेड सामायिक करतात तेव्हा स्विच लागू होतो; इतर आकार सामान्य मार्गावर राहतात. एक आब्लेशन हा स्विच अक्षम करतो आणि नेहमी सामान्य मार्ग वापरतो. निश्चित-ब्लॉक मार्ग सक्षम केल्याने 32K वर 7.7%, 64K वर 27.4%, आणि 128K वर 40.2% ने एंड-टू-एंड डिकोड सुधारले.

पुढील ऑप्टिमायझेशनच्या मर्यादा

काही बदलांमुळे एकाकी ऑपरेशन सुधारले परंतु एंड-टू-एंड इन्फरन्स सुधारला नाही.

अनुमानित डिकोड, जे संपूर्ण मॉडेल सत्यापित करण्यासाठी टोकन्स प्रस्तावित करण्यासाठी लहान मॉडेल वापरते, त्याने बॅच-1 डिकोड 18% वेगाने कमी केले. पडताळणीने दोन ते पाच ओळींच्या गटांवर प्रक्रिया केली, या हार्डवेअरसाठी अकार्यक्षम आकार, आणि ओळींनी अनेकदा वेगवेगळे तज्ज्ञ निवडले, ज्यामुळे वाचलेल्या तज्ज्ञ-वजन डेटाचे प्रमाण वाढले. ड्राफ्टर्सचे आउटपुट व्होकॅब्युलरी कमी केल्याने ड्राफ्टर थ्रूपुट 4.7–5.1% सुधारला, परंतु पूर्ण अनुमानित लूप वेगवान झाला नाही. हा निकाल वर्कलोड-विशिष्ट आहे: ब्लॅकवेलवरील आमचे बॅच केलेले Qwen डिप्लॉयमेंट वेगवेगळ्या परिस्थितींमध्ये अनुमानित डिकोड वापरते.

इतर प्रयोगांमध्ये GPU लॉन्च कमी करणे, संपूर्ण टप्पे ओव्हरलॅप करणे, मोठ्या प्रिफिल टाइल्स वापरणे, व्यापक संलयन लागू करणे, राउटर वेगवान करणे आणि आउटपुट प्रोजेक्शन टोकन निवडीसह एकत्र करणे समाविष्ट होते. कशामुळेही संपूर्ण इन्फरन्स लूप सुधारला नाही.

हार्डवेअर मर्यादांच्या मोजमापांनी मुख्य प्रिफिल आणि डिकोड ऑपरेशन्समध्ये थोडीच उरलेली हेडरुम दर्शवली. MoE GEMM आणि GEMV त्यांच्या प्रवेश नमुन्यांसाठी सर्वात जलद सतत वजन-वाचन दरांच्या 97.9% आणि 90.3% पर्यंत पोहोचले. विरल GEMV मधून अंकगणित काढून टाकल्याने थ्रूपुट केवळ 0.2% बदलला, ज्यामुळे संगणनाऐवजी वजन वाचणे ही मर्यादा घालणारी संसाधन असल्याची पुष्टी झाली. प्रिफिलच्या मॅट्रिक्स गुणाकाराने अशाच प्रकारे अलगावमध्ये सैद्धांतिक मॅट्रिक्स मर्यादेच्या 93% आणि चाचणी केलेल्या मॉडेलच्या आत 80-86% गाठले.

एंड-टू-एंड कार्यक्षमता

एंड-टू-एंड तुलना दोन्ही इंजिनांमध्ये समान 4-बिट चेकपॉइंट बाइट्स लोड करते आणि एका 40-कोर, 128 GB M5 Max वर एका वेळी एक विनंती चालवते. प्रत्येक फेरीमध्ये, पार्श्वभूमी लोड आणि चिप तापमानातील बदलांमधील पूर्वग्रह कमी करण्यासाठी दोन इंजिने पर्यायी क्रमाने चालतात. आम्ही MLX-LM च्या सर्वात वेगवान थेट-निर्मिती मार्गाशी तुलना करतो, त्याच्या सर्व्हरशी नाही, त्यामुळे मोजमाप सर्व्हिंग ओव्हरहेडऐवजी मॉडेल अंमलबजावणीवर केंद्रित असते.

स्वूपमध्ये प्रिफिलसाठी दहा प्रॉम्प्ट लांबी आणि डिकोडसाठी दहा संदर्भ लांबी समाविष्ट आहेत, 256 ते 128K टोकन्सपर्यंत. प्रिफिल थ्रूपुट प्रथम वाढतो कारण इंजिन अधिक टोकन्सवर निश्चित सेटअप खर्च पसरवते. प्रिफिल सुमारे 4K-टोकन प्रॉम्प्टवर शिखर गाठतो, नंतर खाली पडतो कारण जसजसा प्रॉम्प्ट वाढतो तसतसे दहा फुल-अटेंशन लेअर्स अधिक कार्य करतात. लहान संदर्भांवर डिकोड जवळपास सपाट राहते आणि वाढती KV कॅशे वाचणे महत्त्वपूर्ण झाल्यावर कमी होते. सानुकूल इंजिन प्रत्येक रेकॉर्ड केलेल्या लांबीवर वेगवान आहे.

विशेषीकृत अंमलबजावणीमुळे फ्लोटिंग-पॉइंट ऑपरेशन्सचा क्रम बदलू शकत असल्याने, आम्ही MLX-LM विरुद्ध संख्यात्मक सुसंगतता देखील तपासली. शिक्षक-सक्तीच्या तुलनेत, दोन्ही इंजिनांनी 192 स्थानांपैकी प्रत्येक ठिकाणी समान संदर्भ उपसर्गावरून पुढील टोकनचा अंदाज लावला, ज्यामुळे पूर्वीच्या फरकांना नंतरच्या इनपुटवर परिणाम करण्यापासून प्रतिबंधित केले. लिलीची पर्प्लेक्सिटी केवळ 0.04% जास्त होती, आणि तिने चाचणी केलेल्या स्थानांच्या 96.35% वर समान शीर्ष-क्रमांकित टोकन निवडले.

स्थानिक प्लॅटफॉर्मसाठी तयार केलेले

Apple सिलिकॉन हे लहान डॅटाम्झर GPU नाही. हे स्वतःच्या हार्डवेअर आणि सॉफ्टवेअर वैशिष्ट्यांसह एक संपूर्ण स्थानिक इन्फरन्स प्लॅटफॉर्म आहे. युनिफाइड मेमरी एकाच नोडला मॉडेल आणि राज्य किती धरून ठेवू शकते यावर खूप उच्च मर्यादा देते. M5 न्यूरल अ‍ॅक्सिलेटर प्रिफिलमधील दाट मॅट्रिक्स कार्य शोषून घेतात. वेक्टर ALU बँडविड्थ-बाउंड, कमी-पुनरावृत्ती उरलेला भाग डिकोडमध्ये हाताळतात.

Qwen विशेषीकरणासाठी पुढील संधी जोडतो: तज्ज्ञ रुटिंग आणि पुनरावृत्ती स्थिती GPU वर ठेवा, अनावश्यक मध्यवर्ती नष्ट करा, स्वतंत्र कार्य ओव्हरलॅप करा, सामायिक KV डेटा पुन्हा वापरा, आणि वर्कलोड आकारानुसार कर्नल अनुकूल करा.

मॉडेल- आणि प्लॅटफॉर्म-विशिष्ट ऑप्टिमायझेशनसह, एक मॅक मोठ्या विरल मॉडेलला कार्यक्षमतेने चालवू शकतो. भविष्यातील काम मॉडेल्स, चिप्स आणि सर्व्हिंग वर्कलोड्समध्ये कव्हरेज विस्तृत करेल आणि एका कॉन्फिगरेशनवर येथे सत्यापित केलेल्या यंत्रणांना अधिक सामान्य रनटाइम धोरणात बदलेल.

व्यापक तत्त्व म्हणजे इंजिनला मॉडेलच्या वास्तुकलेशी आणि हार्डवेअरच्या विशिष्ट संगणन आणि मेमरी मार्गांशी जुळवणे. जसे सीमा ओपन-वेट मॉडेल्स आणि हार्डवेअर विकसित होतात, उच्च-कार्यक्षमता स्थानिक इन्फरन्स वाढत्या प्रमाणात त्यांच्यातील फरक अमूर्त करणाऱ्याऐवजी दोन्हीनुसार तयार केलेल्या इंजिनवर अवलंबून असेल.