The Other Half of the Memory Wall: Serving 35B MoEs from SSD with Trained Routing Prediction

Yu Lin; Yiming Wang; Runyuan Cai; Hanze Liu; Xiaodong Zeng

शोध पत्र · arXiv 2609.18063

मेमोरी वॉल का दूसरा पहलू: ट्रेन्ड रूटिंग प्रेडिक्शन के साथ SSD से 35B MoE की सर्विंग

The Other Half of the Memory Wall: Serving 35B MoEs from SSD with Trained Routing Prediction

Yu Lin, Yiming Wang, Runyuan Cai, Hanze Liu, Xiaodong Zeng

प्रकाशित 2026-09-16cs.AI arXiv PDF

अनुवाद और AI के साथ पढ़ें

तीन पंक्तियों में सारांश

🎯 35B-क्लास MoE को int4 पर भी ~19.5 GB की आवश्यकता होती है, जो कंज्यूमर RAM से अधिक है; वहीं सामान्य SSD ऑफलोडिंग डिस्क लेटेंसी पर अटक जाती है क्योंकि लेयर रूटिंग पिछली लेयर के आउटपुट पर निर्भर करती है।
🔧 Edge0 प्रति-लेयर prerouter का उपयोग करके SSD से MoE एक्सपर्ट्स को स्ट्रीम करता है, जो 'prediction as routing' को अपनाते हुए एक टोकन पहले रूटिंग की भविष्यवाणी करता है, जिसे स्टूडेंट पाथ पर डिस्टिल्ड अनमर्ज्ड पैरेलल recovery LoRA के साथ जोड़ा गया है।
📈 24 GB Mac mini M4 Pro पर 2.9 GiB पीक मेमोरी के भीतर एक 35B MoE को ~20.4 tok/s पर सर्व करता है, जो fp16 टीचर से औसतन केवल 3.9 बेंचमार्क अंक पीछे है।

सार

कंज्यूमर हार्डवेयर पर Mixture-of-experts (MoE) का इन्फरेंस वेट मेमोरी द्वारा सीमित होता है: 4-बिट पर 35B-श्रेणी का मॉडल 19.5GB का होता है, और स्पार्सिटी प्रति टोकन कंप्यूट को कम करती है, न कि उन बाइट्स को जिन्हें मेमोरी में रखना आवश्यक होता है। केवल SSD पर सामान्य ऑफलोडिंग करने से कोई मदद नहीं मिलती, क्योंकि लेयर N+1 के एक्सपर्ट्स को लेयर N का आउटपुट आने से पहले चुना जाना आवश्यक होता है, जिससे री-ड्स कंप्यूट के पीछे छिपने लायक जल्दी शुरू नहीं हो सकते। हम Edge0 प्रस्तुत करते हैं, जो एक स्ट्रीमिंग MoE इन्फरेंस इंजन है और एक prerouter के माध्यम से इस अंतर को समाप्त करता है: एक प्रति-लेयर हेड अगली लेयर की रूटिंग की भविष्यवाणी एक टोकन पहले ही कर लेता है, और इस प्रेडिक्शन को ही वास्तविक रूटिंग मान लिया जाता है, जिससे स्टेज किया गया एक्सपर्ट सेट रूट किए गए सेट के समान होता है और कुछ भी ड्रॉप नहीं होता है। स्टूडेंट पाथ पर प्रशिक्षित एक unmerged recovery LoRA, int4 क्वांटाइजेशन और रूटिंग रिप्लेसमेंट के कारण कम हुई गुणवत्ता की भरपाई करता है। 24GB वाली एक अकेली मशीन पर, Edge0 पाँच पब्लिक बेंचमार्क के औसत पर अपने fp16 टीचर के कुछ ही अंकों के भीतर रहते हुए, 3GiB के पीक एक्टिव मेमोरी में 35B MoE को 20tok/s पर सर्व करता है। एक 8B टीयर भी इसी फ्रेमवर्क पर चलता है, और यह फ्रेमवर्क, चेकपॉइंट्स व एडेप्टर्स ओपन सोर्स हैं।

मूल सार (अंग्रेज़ी)

Mixture-of-experts (MoE) inference on consumer hardware is bounded by weight memory: a 35B-class model is 19.5GB at 4-bit, and sparsity shrinks the compute per token, not the bytes that must be held. Naive offloading to SSD does not help on its own, because layer N+1's experts must be chosen before layer N's output exists, so the reads cannot start early enough to hide behind compute. We present Edge0, a streaming MoE inference engine that closes the gap with a prerouter: a per-layer head predicts the next layer's routing one token ahead, and the prediction is consumed as the routing itself, so the staged expert set equals the routed set and nothing is dropped. An unmerged recovery LoRA, trained on the student path, pays back the quality lost to int4 quantization and routing replacement. On a single 24GB machine, Edge0 serves a 35B MoE at 20tok/s inside 3GiB of peak active memory, within a few points of its fp16 teacher on average across five public benchmarks. An 8B tier runs on the same framework, and the framework, checkpoints, and adapters are open source.

पढ़ने की गाइड

समस्या और प्रेरणा

Mixture-of-Experts (MoE) मॉडल प्रति टोकन एक्टिव पैरामीटर्स को कम रखते हुए (उदा. 35B मॉडल में प्रति टोकन ~3B पैरामीटर्स को एक्टिव करना) कुल पैरामीटर्स को स्केल करते हैं। हालाँकि, कंज्यूमर हार्डवेयर पर इन्फरेंस स्टेटिक वेट मेमोरी द्वारा बाधित रहता है: int4 में क्वांटाइज करने पर भी 35B MoE को लगभग 19.5 GB की आवश्यकता होती है। 24 GB की मशीन पर, यह ऑपरेटिंग सिस्टम और KV कैश के लिए जगह नहीं छोड़ता, जिससे गंभीर मेमोरी पेजिंग होती है और डिकोड थ्रूपुट घट जाता है (उदा. M4 Pro Mac mini पर रेजिडेंट मॉडल के लिए 3.9 tok/s)।

एक्सपर्ट वेट्स को SSD पर ऑफलोड करने से मेमोरी फुटप्रिंट कम होता है, लेकिन सामान्य ऑन-डिमांड लोडिंग एक गंभीर सीरियलाइजेशन बॉटलनेक पैदा करती है: लेयर $N + 1$ के लिए एक्सपर्ट का चयन लेयर $N$ के आउटपुट पर निर्भर करता है। SSD रीड्स कंप्यूट के पीछे छिपने के लिए पर्याप्त जल्दी शुरू नहीं हो सकते, जिससे प्रति टोकन प्रति लेयर डिस्क लेटेंसी पर एग्जीक्यूशन रुक जाता है।

मुख्य विचार

वेट टीयर के रूप में SSD: MoE एक्सपर्ट वेट्स SSD पर int4 फॉर्मेट में रहते हैं और एक कॉम्पैक्ट स्ट्रीमिंग पूल में मेमोरी-मैप (mmap) किए जाते हैं। पीक मेमोरी कुल मॉडल पैरामीटर आकार के बजाय एक्टिव वर्किंग सेट द्वारा निर्धारित होती है।
रूटिंग के रूप में प्रेडिक्शन (Prediction as Routing): लेयर $N$ के पूरा होने की प्रतीक्षा करने के बजाय, एक प्रति-लेयर prerouter पूरे एक टोकन पहले ही लेयर $N + 1$ की रूटिंग का अनुमान लगा लेता है। महत्वपूर्ण बात यह है कि मॉडल डिकोड के समय प्रेडिक्शन को ही वास्तविक रूटिंग मानता है। यह डिज़ाइन के स्तर पर ही ड्रॉप हुए टोकन या फॉलबैक लोड को समाप्त कर देता है क्योंकि स्टेज किया गया सेट रूट किए गए सेट के बिल्कुल समान होता है।
अनमर्ज्ड रिकवरी LoRA (Unmerged Recovery LoRA): 4-बिट क्वांटाइजेशन और अनुमानित रूटिंग के कारण गुणवत्ता का नुकसान होता है। Edge0 डिस्टिलेशन के माध्यम से स्टूडेंट रूटिंग पाथ पर एक लो-रैंक एडेप्टर को प्रशिक्षित करता है। चूंकि LoRA को int4 बेस में मर्ज करने और फिर से क्वांटाइज करने से अधिकांश डेल्टा मिट जाता है (लॉगिट्स स्तर पर केवल 18% ही बच पाता है), Edge0 इस एडेप्टर को नगण्य ओवरहेड पर एक पैरेलल अनमर्ज्ड डेल्टा के रूप में सर्व करता है।

विधि

Prerouter हेड आर्किटेक्चर

लेयर $N$ का प्रत्येक prerouter हेड टोकन $t$ पर पोस्ट-अटेंशन लेयर नॉर्म आउटपुट पर चलता है और टोकन $t + 1$ पर लेयर $N + 1$ के लिए रूटिंग की भविष्यवाणी करता है। इनपुट फीचर्स हिडन स्टेट को दो top- $k$ वन-हॉट रूटिंग वेक्टर्स (वर्तमान और पिछला टोकन) के साथ जोड़ते हैं:

$dim = d_{model} + 2 E$

35B टीयर के लिए, यह $2048 + 2 \times 256 = 2560$ है। यह हेड गणना करता है:

$head (x) = fc_{2} (erf - gelu (fc_{1} (x))) + ℓ (x)$

जहाँ $ℓ (x)$ अगली लेयर के राउटर वेट्स से वार्म-स्टार्ट किया गया एक लीनियर रेजिड्यूअल पाथ है।

अनमर्ज्ड एडेप्टर एग्जीक्यूशन

सर्व की जाने वाली MoE लेयर क्वांटाइज्ड बेस बाइट्स को बनाए रखती है और पैरेलल LoRA डेल्टा चलाती है:

$y = W_{int4} (x) + \frac{α}{r} B A x$

तीन-चरणीय प्रशिक्षण (Three-Phase Training)

हेड डिस्टिलेशन: Prerouter हेड्स को बाकी नेटवर्क को फ्रीज रखते हुए टीचर के अगली लेयर के वास्तविक राउटर निर्णयों की नकल करने के लिए प्रशिक्षित किया जाता है।
स्टूडेंट-पाथ SFT: रिकवरी LoRA एडेप्टर्स (जो अटेंशन, लीनियर-अटेंशन, और शेयर्ड-एक्सपर्ट प्रोजेक्शन से जुड़े हैं, लेकिन स्ट्रीमिंग एक्सपर्ट्स से नहीं) को prerouter के माध्यम से रूटिंग को बाध्य करते हुए टीचर-जनरेटेड टेक्स्ट का उपयोग करके प्रशिक्षित किया जाता है।
ऑन-पॉलिसी डिस्टिलेशन: स्टूडेंट प्री-रूटेड एग्जीक्यूशन के तहत टेक्स्ट उत्पन्न करता है, fp16 बेस टीचर इन जनरेशन्स को स्कोर करता है, और टेल-अवेयर top- $k$ प्रोबेबिलिटी ऑब्जेक्टिव के साथ रिवर्स Kullback–Leibler डाइवर्जेंस के माध्यम से LoRA पैरामीटर्स को ऑप्टिमाइज़ किया जाता है।

परिणाम

थ्रूपुट और मेमोरी (Mac mini M4 Pro, 24 GB):
- edge0-35b ( $K = 4$ ): 19.5 GB int4 चेकपॉइंट से केवल 2.9 GiB पीक एक्टिव मेमोरी का उपयोग करके 20.4 tok/s डिकोड प्राप्त करता है (पूरी तरह से मेमोरी में रहने पर 3.9 tok/s और 18.2 GiB की तुलना में)।
- edge0-8b ( $K = 8$ ): 1.5 GiB पीक मेमोरी के भीतर 28.0 tok/s डिकोड प्राप्त करता है।
प्रीफेचिंग स्पीडअप (MacBook M2, 16 GB):
- विशुद्ध रूप से रिएक्टिव ऑन-डिमांड स्ट्रीमिंग की तुलना में, 1-टोकन prerouter प्रीफेच सभी रूटिंग चौड़ाई में डिकोड स्पीड को +80% से +84% तक बढ़ाता है ( $K = 2$ : 4.8 से 8.6 tok/s; $K = 4$ : 3.5 से 6.4 tok/s; $K = 8$ : 1.8 से 3.3 tok/s)।
- मेन-थ्रेड ब्लॉक लोड समय काफी कम हो जाता है (उदा. $K = 4$ पर 244.0 ms से घटकर 101.9 ms प्रति स्टेप)।
मॉडल गुणवत्ता (OpenCompass):
- पाँच बेंचमार्क (AIME 2026, HumanEval, GPQA-Diamond, MMLU-Pro, IFBench) में, edge0-35b का औसत स्कोर fp16 Qwen3.6 बेस के 83.2 के मुकाबले 79.2 रहा (3.9-अंकों का अंतर)।
- edge0-8b ने Ling 3.0 tiny fp16 बेस के 72.7 के मुकाबले 69.9 प्राप्त किया (2.8-अंकों का अंतर)।

सीमाएं और खुले प्रश्न

सर्विंग कॉनकरेंसी: Edge0 सिंगल रिक्वेस्ट के लिए FIFO-क्रमबद्ध संचालित होता है; बैचिंग एक्सपर्ट वर्किंग सेट और कैशिंग डायनामिक्स को ऐसे तरीकों से बदलती है जिन्हें सिंगल-स्ट्रीम प्रोफाइल में संबोधित नहीं किया गया है।
CPU ग्राफ-बिल्डिंग ओवरहेड: डिकोड स्पीड CPU-साइड फ्रेमवर्क ओवरहेड (40 लेयर्स में MLX ग्राफ बिल्डिंग में ~44 ms/स्टेप) द्वारा सीमित है, जिसे SSD ऑप्टिमाइज़ेशन कम नहीं कर सकते।
रीजनिंग पर गुणवत्ता में गिरावट: प्रदर्शन में गिरावट मुख्य रूप से जटिल रीजनिंग बेंचमार्क पर केंद्रित है (उदा. 35B के लिए AIME पर 6.1 अंकों की गिरावट और 8B के लिए 10.0 अंकों की गिरावट)।
हार्डवेयर निर्भरता: Prerouter का सापेक्ष थ्रूपुट लाभ डिस्क लेटेंसी पर निर्भर करता है और वेट कैश वार्म होने पर या अत्यधिक तेज़ स्टोरेज टीयर पर चलने पर कम हो जाता है।

मुख्य शब्द

Edge0: एक स्ट्रीमिंग MoE इन्फरेंस इंजन जो रूटिंग प्रेडिक्शन के जरिए कंप्यूट के साथ रीड्स को ओवरलैप करके SSD से मॉडल सर्व करता है और अनमर्ज्ड रिकवरी LoRA के साथ गुणवत्ता के नुकसान की भरपाई करता है।
prerouter: एक हल्का प्रति-लेयर न्यूरल हेड जो लेयर N के हिडन स्टेट और हालिया रूटिंग हिस्ट्री का उपयोग करके एक टोकन पहले ही लेयर N+1 के रूटिंग निर्णयों का अनुमान लगाता है।
prediction-as-routing: यह डिज़ाइन सिद्धांत जहाँ prerouter के अनुमानित एक्सपर्ट सेट को डिकोड समय पर वास्तविक लेयर रूटिंग के रूप में सीधे निष्पादित किया जाता है, जिससे रनटाइम टोकन ड्रॉपिंग या मिस पेनल्टी समाप्त हो जाती है।
recovery LoRA: स्टूडेंट प्री-रूटेड पाथ पर डिस्टिलेशन के माध्यम से प्रशिक्षित एक लो-रैंक एडेप्टर, जिसे int4 क्वांटाइजेशन और रूटिंग एप्रोक्सीमेशन त्रुटियों की भरपाई के लिए एक अनमर्ज्ड पैरेलल ब्रांच के रूप में चलाया जाता है।
staged decode: फिक्स्ड-स्लॉट डबल बफरिंग का उपयोग करने वाला एक एक्सपर्ट एग्जीक्यूशन मोड, जहाँ प्रति-स्टेप एलोकेशन से बचने के लिए GPU पर स्लॉट टेबल के माध्यम से एक्सपर्ट इंडेक्स मैप किए जाते हैं।
incremental sticky-slot stack (incr_stack): एक मैकेनिज्म जो सभी एक्सपर्ट टेंसरों को फिर से जोड़ने के बजाय, इन-प्लेस परसिस्टेंट स्टैक्ड एक्सपर्ट टेंसरों को अपडेट करता है और प्रति स्टेप केवल बदले हुए एक्सपर्ट स्लॉट को रीराइट करता है।
unmerged delta: LoRA एडेप्टर को क्वांटाइज्ड बेस वेट्स में फ्यूज करने के बजाय उनके समानांतर निष्पादित करना, जिससे री-क्वांटाइजेशन के दौरान फाइन-ग्रेन्ड एडेप्टेशन का नुकसान नहीं होता।
memory wall: प्रोसेसर थ्रूपुट और DRAM बैंडविड्थ/कैपेसिटी के बीच का अंतर, जो लोकल LLM सर्विंग में डायनामिक KV-कैश स्टेट और स्टेटिक मॉडल वेट स्टोरेज के बीच विभाजित होता है।

प्रश्न

Edge0 प्रेडिक्शन को हिंट मानने के बजाय रूटिंग के रूप में प्रेडिक्शन (prediction as routing) का उपयोग क्यों करता है?

प्रेडिक्शन को हिंट मानने पर रनटाइम फॉलबैक या प्रेडिक्शन मिस होने पर टोकन ड्रॉप करने की आवश्यकता होती है। प्रेडिक्शन को सीधे रूटिंग के रूप में उपयोग करने से यह सुनिश्चित होता है कि स्टेज किया गया एक्सपर्ट सेट हमेशा रूट किए गए सेट से मेल खाता है, जिससे एप्रोक्सीमेशन की लागत ऑफलाइन ट्रेनिंग में स्थानांतरित हो जाती है जहाँ डिस्टिलेशन गुणवत्ता को पुनर्प्राप्त कर लेता है।

रिकवरी LoRA को 4-बिट वेट्स में फ्यूज करने के बजाय अनमर्ज्ड क्यों सर्व किया जाता है?

LoRA वेट डेल्टा का RMS लगभग 10^-3 होता है, जो 4-बिट क्वांटाइजेशन के स्टेप साइज से छोटा है। मर्ज करने और दोबारा क्वांटाइज करने से एडेप्टर का अधिकांश प्रभाव समाप्त हो जाता है, जिससे लॉगिट्स स्तर पर केवल 18% प्रभाव ही बचता है, जबकि एक पैरेलल अनमर्ज्ड डेल्टा 100% सटीकता बनाए रखता है।

Edge0 उसी टोकन के भीतर अनुमान लगाने के बजाय एक टोकन आगे का अनुमान क्यों लगाता है?

उसी टोकन के भीतर प्रेडिक्शन करने के लिए प्रति-लेयर GPU पाइपलाइन फ्लश की आवश्यकता होती है जिसमें प्रति स्टेप 30 से 100 ms लगते हैं—जो उस लोड समय से अधिक है जिसे यह छिपाने का प्रयास करता है। एक पूरे टोकन से आगे रहने पर सभी स्टेज की गई लेयर्स को एक ही बैच फ्लश में प्रेडिक्ट किया जा सकता है, जिससे SSD I/O को छिपाने के लिए कंप्यूट टाइम का एक पूरा डिकोड स्टेप मिल जाता है।

24 GB मशीन पर पूरी तरह से इन-मेमोरी बेसलाइन की तुलना में Edge0 तेज़ी से डिकोड कैसे करता है?

पूरी तरह से मेमोरी में स्थित 19.5 GB का मॉडल 18.2 GiB RAM की खपत करता है, जिससे ऑपरेटिंग सिस्टम और KV कैश के लिए लगभग कोई जगह नहीं बचती। इसके कारण गंभीर OS मेमोरी पेजिंग होती है और डिकोड गिरकर 3.9 tok/s पर आ जाता है। Edge0 ~2.9 GiB पीक एक्टिव मेमोरी में SSD से एक्सपर्ट्स को स्ट्रीम करता है, जिससे OS मेमोरी प्रेशर पूरी तरह से बच जाता है और यह 20.4 tok/s पर चलता है।

इस पृष्ठ की व्याख्या AI ने शोध पत्र से लिखी है और इसमें गलतियाँ हो सकती हैं। विवरण के लिए मूल शोध पत्र पढ़ें।

अनुवाद और AI के साथ पढ़ें