LLM को MegaKernel में compile करके low-latency inference हासिल करना
(zhihaojia.medium.com)- LLM inference को एकल megakernel में अपने-आप बदलने वाला compiler विकसित किया गया है
- MegaKernel(Persistent kernel) तरीका LLM inference में computation और communication को पूरी तरह एक GPU kernel में एकीकृत करता है, जिससे बहुत कम latency संभव होती है
- मौजूदा ML framework या kernel library की बिखरी हुई संरचना के कारण पूरी pipeline को एक single kernel में बदलना बहुत कठिन है
- Mirage Persistent Kernel(MPK) compiler और runtime system के माध्यम से multi-GPU LLM inference को अपने-आप high-performance megakernel में बदलता है
- MPK computation graph को fine-grained task graph में बदलता है, जिससे software pipelining और computation-communication overlap को अधिकतम किया जा सके
- MPK लागू करने पर मौजूदा systems की तुलना में token generation latency घटती है, और GPU की संख्या बढ़ने पर performance improvement और अधिक बढ़ता है
अवलोकन और MegaKernel तरीके के फायदे
- बड़े language model(LLM) inference में latency घटाने का एक प्रभावी तरीका यह है कि सभी computation और communication प्रक्रियाओं को एक single megakernel(सुसंगत kernel) में fuse कर दिया जाए
- इस तरीके में पूरे model की layer-wise operations और GPU के बीच communication सहित सभी processing को एक ही GPU kernel बिना रुकावट के पूरा करता है
- इसके मुख्य फायदे इस प्रकार हैं
- बार-बार होने वाली kernel calls को हटाकर kernel launch overhead समाप्त करना
- layers के पार software pipelining लागू करना
- computation और communication को साथ-साथ चलाकर latency छिपाना
मौजूदा सीमाएँ और MPK का आगमन
- मौजूदा PyTorch, Triton, TVM जैसे ML frameworks स्वाभाविक रूप से end-to-end megakernel auto-generation को support नहीं करते
- वास्तविक LLM systems NCCL/NVSHMEM(communication), FlashInfer/FlashAttention(attention), CUDA/Triton(custom operations) जैसी कई kernel libraries के संयोजन से बने होते हैं, इसलिए उन्हें एक single kernel में एकीकृत करना कठिन है
- इसी पृष्ठभूमि में CMU, UW, Berkeley, NVIDIA और Tsinghua के शोधकर्ताओं ने Mirage Persistent Kernel(MPK) विकसित किया
- MPK compiler और runtime को जोड़कर पूरे LLM inference pipeline को अपने-आप high-performance megakernel में बदल देता है
MPK का मुख्य मूल्य
- MPK kernel launch overhead को पूरी तरह हटाता है और layers के बीच computation/data loading/communication overlap को अधिकतम करके ultra-low-latency LLM inference environment संभव बनाता है
- वास्तविक test (39-token prompt, 512 token generation, speculative decoding का उपयोग नहीं) में,
- single NVIDIA A100 40GB GPU environment में vLLM/SGLang जैसे मौजूदा optimized systems की token decoding latency (14.5ms) की तुलना में MPK इसे 12.5ms तक घटा देता है
- यह मान theoretical lower bound(10ms) के करीब है (1.6TB/s memory bandwidth, 16GB weight loading के आधार पर)
- multi-GPU environment में computation और communication को पूरी तरह integrate करने के कारण, GPU की संख्या बढ़ने पर MPK की performance बढ़त और अधिक स्पष्ट हो जाती है
MPK की कार्यप्रणाली का विस्तृत ढाँचा
Part 1. Compiler – LLM computation graph → task graph रूपांतरण
- सामान्यतः LLM computation को computation graph के रूप में दर्शाया जाता है, जहाँ प्रत्येक operation (जैसे matrix multiplication, attention) या communication operation (जैसे all-reduce) एक node होता है और data dependency edge के रूप में व्यक्त होती है
- पारंपरिक design में प्रति operator अलग kernel चलाना आम बात है, लेकिन यह वास्तविक data unit की dependency के बजाय केवल kernel-level dependency को दर्शाता है, इसलिए pipelining के अवसर सीमित हो जाते हैं
- उदाहरण: यदि matrix multiplication के बाद allreduce हो, तो पारंपरिक रूप से पूरा matrix multiplication खत्म होने के बाद ही allreduce शुरू होता है। जबकि वास्तव में data को टुकड़ों में बाँटकर partial execution और dependency का उपयोग किया जा सकता है
- MPK compiler computation graph को अधिक सूक्ष्म स्तर पर तोड़कर, वास्तविक data units के अनुरूप fine-grained task graph में अपने-आप बदल देता है
- प्रत्येक task (आयत) एक individual GPU SM को सौंपा गया computation/communication unit होता है
- प्रत्येक event (वृत्त) tasks के बीच synchronization point होता है
- task और event के बीच edges के जरिए data/control dependency को कुशलता से व्यक्त किया जाता है
- इस task graph की वजह से MPK computation और communication को आंशिक रूप से या parallel रूप में अधिक overlap करने देता है
- Mirage kernel superoptimizer प्रत्येक task के लिए उपयुक्त high-performance CUDA implementation भी अपने-आप generate करता है
Part 2. Runtime – megakernel के भीतर task graph का execution
- MPK runtime task graph को GPU के एक ही kernel (megakernel) के भीतर पूरी तरह execute करता है
- GPU के सभी SM(Streaming Multiprocessors) को static रूप से worker और scheduler भूमिकाओं में बाँटा जाता है
Worker
- प्रत्येक worker SM unit पर काम करता है और अपनी dedicated task queue संभालता है
- loop तरीके से
- queue से अगला task लाना
- उसे execute करना (जैसे matmul, attention, data transfer)
- पूरा होने पर event को सूचित करना
- फिर यही प्रक्रिया दोहराना
- इससे प्रत्येक worker के resource utilization का optimization होता है और asynchronous layer operations संभव होते हैं
Scheduler
- distributed scheduler प्रत्येक SM में single warp unit पर काम करता है, और अधिकतम 4 schedulers एक साथ चल सकते हैं
- प्रत्येक scheduler activated event queue को manage करता है और जिन tasks की conditions पूरी हो चुकी हों उन्हें workers को assign करता है
- इससे centralized synchronization overhead के बिना बड़े पैमाने पर distributed task processing संभव होती है
Event-आधारित execution तरीका
- task पूरा होने पर वह specific event counter बढ़ाता है। जब counter threshold तक पहुँचता है, event activate हो जाता है और scheduler queue में डाल दिया जाता है
- scheduler उस event पर dependent अगले tasks को execute कराता है
- इसकी वजह से fine-grained software pipelining और computation-communication overlap स्वाभाविक रूप से संभव हो जाते हैं
- उदाहरण: एक layer का matmul और दूसरी layer का attention एक साथ चल सकते हैं
- matmul का आंशिक result मिलते ही allreduce communication शुरू किया जा सकता है
- क्योंकि पूरा scheduling और task switching single kernel context के भीतर होता है, tasks के बीच overhead 1–2 microsecond(μs) स्तर तक बहुत कम रहता है
भविष्य की दिशा
-
MPK का लक्ष्य: developers केवल थोड़े से Python code (लगभग कुछ दर्जन lines) लिखकर आसानी से LLM को megakernel में compile कर सकें और अधिकतम performance हासिल कर सकें
-
प्रमुख विकास दिशाएँ
- नवीनतम GPU architecture support: उदाहरण के लिए NVIDIA Blackwell, warp-level specialization जैसी विधियाँ
- dynamic workload handling: mixture-of-experts(MoE) जैसे dynamic control flow वाले models के लिए compilation strategy पर शोध
- advanced task scheduling: priority-based, throughput optimization जैसी आधुनिक policies पर शोध और उनके उपयोग की संभावना
-
MPK GPU-आधारित LLM inference workloads के compile और execution तरीके में एक बुनियादी मोड़ प्रस्तुत करता है, और समुदाय के साथ सहयोग के विस्तार की अपेक्षा रखता है
अतिरिक्त सामग्री
- MPK(Mirage Persistent Kernel) code, documentation और नवीनतम research results GitHub(https://github.com/mirage-project/mirage) पर देखे जा सकते हैं
1 टिप्पणियां
Hacker News राय
लेखक के लिए, on-GPU interpreter तरीका एक बहुत आशाजनक भविष्य की दिशा लगता है, यह काफ़ी दिलचस्प है। लगभग इसी तरह का अप्रोच दिखाने वाला एक और शोध भी है, इसलिए संबंधित पोस्ट देखने की सलाह है। CUDA का मूल programming model (जैसे kernel launch) सूक्ष्म task-based parallelism के लिए बायपास किया जा रहा है, और इस तरीके से hardware utilization वास्तव में बढ़ता हुआ दिख रहा है। सोचने वाली बात है कि क्या CUDA हमें कई तरह से सीमित करके रखे हुए था। यह उम्मीद भी है कि लेखक का शोध PyTorch के experimental backend में आ सके। और, पहले हिस्से के दो पैराग्राफ लगभग एक जैसे हैं, इसलिए एक छोटा typo भी बताना चाहूँगा.
मैंने कुछ समय तक vLLM और SGLang में काफ़ी क़रीब से काम किया है, और मुझे पूरा यक़ीन है कि यह प्रोजेक्ट ठीक उसी का आदर्श अगला चरण है। computation dependency graph का analysis, operations को fuse करना, और tasks को अधिक समझदारी से schedule करना — ये सब बहुत प्रभावशाली लगा। टीम को बधाई.
मैंने पोस्ट और github README देखा, और यह सच में बहुत शानदार प्रोजेक्ट लगा। जानना चाहता हूँ कि क्या इस तरह की optimization सिर्फ inference ही नहीं बल्कि training stage पर भी लागू हो सकती है। खासकर यह समझता हूँ कि backward operations और gradient communication का fusion एक चुनौती है। जहाँ तक मुझे पता है, अभी dynamic workloads (जैसे MoE) supported नहीं हैं, लेकिन हाल की एक paper FlashDMoE: Fast Distributed MoE in a Single Kernel है जो एक ही kernel में MoE को संभालती है.
पोस्ट और README तक पढ़ने के लिए धन्यवाद। training stage को support करना संभव है, लेकिन आम तौर पर training kernels बड़े होते हैं, इसलिए kernel launch overhead वहाँ उतनी बड़ी समस्या नहीं बनती; इसी कारण inference, खासकर low-latency inference, इसका बड़ा लाभार्थी है। आपने जो FlashDMoE paper साझा की वह भी बहुत रोचक लगी, और MoE model support को अगला लक्ष्य बनाया गया है।
व्यक्तिगत रूप से, gradient-based learning optimization पर समय लगाना मुझे कुछ हद तक संदिग्ध लगता है। व्यवहार में कई training tasks में discrete values की प्रकृति होती है, और मुझे नहीं लगता कि gradient-based learning उनसे अच्छी तरह निपट पाती है.
अगला कदम तो सीधे Verilog में compile करके aliexpress से अपना LLM hardware खरीदना होना चाहिए, यही सपना है.
Chisel जैसी hardware technologies का परिचय देने वाली पोस्ट साझा कर रहा हूँ। AI और GPU के आने से पहले software से सीधे hardware तक जाने का यह विचार एक काफ़ी आशाजनक दिशा माना जाता था। CPU का विकास ठहराव में है, और software तथा hardware के बीच की परतों को और optimize करने की इच्छा लगातार रही है, लेकिन GPU-style parallel computing आगे भी मुख्य acceleration approach बनी रहने की संभावना अधिक है। general-purpose CPU शायद अंततः GPU को manage करने वाले छोटे brain की भूमिका में रह जाएगा। फिर भी, software से सीधे hardware में जाना मुख्यधारा बन पाएगा, ऐसा नहीं लगता.
5~10 साल बाद अगर LLM की architecture स्थिर हो जाती है, तो उसे सीधे hardware पर map करना व्यावहारिक हो सकता है। मौजूदा तकनीक के साथ भी, सैकड़ों अरब parameters को लगभग 1.5-bit के ultra-low-precision logic gates से एक single wafer में समेटा जा सकता है, ऐसी संभावना है। precision बढ़ने पर gates की संख्या geometric रूप से बढ़ती है, इसलिए अभी के लिए weight memory को बनाए रखना और compute units को share करना अधिक efficient है। भविष्य में ultra-low-precision LLM का विकास एक ज़रूरी कार्य बन सकता है.
training cost पहले से ही बहुत ऊँची है, उस पर mask cost भी जोड़ दी जाए तो हालात और कठिन हो जाएँगे — इस पर एक मज़ाक, और साथ ही यह ठंडी लेकिन यथार्थवादी टिप्पणी कि AI hardware startups काफ़ी समय से इसी दिशा में कोशिश करते आ रहे हैं.
अगर LLM-in-a-box जैसा तरीका वास्तव में मौजूद हो, तो वह काफ़ी आकर्षक होगा। जल्द ही offline (air-gap) environment में काम करने का अवसर मिल सकता है, और ऐसी solution वहाँ बहुत उपयोगी लगेगी.
मैंने खुद Modal GPU environment में कोड चलाकर देखा, और शोध में दावा किए गए performance improvement numbers वास्तव में reproduce हुए। mirage प्रोजेक्ट का result code साझा कर रहा हूँ। Triton + FlashInfer संयोजन में per-token latency लगभग 19.2ms थी, जबकि MPK में उसी condition पर यह 7.7ms तक काफ़ी सुधर गई.
मैंने पहले एक छोटा CUDA contest किया था। वह image या vision domain का parallel algorithm था, और मैंने स्मार्ट दिखने के लिए intermediate results को memory में cache किया। contest के परिणाम आने पर देखा कि दूसरों ने मुझसे कहीं तेज़ code submit किया था, और मैं हैरान रह गया। वजह यह थी कि उन्होंने intermediate results को cache ही नहीं किया, बल्कि बार-बार फिर से compute किया। memory round-trip की तुलना में compute cost बहुत कम थी। मुझे लगता है यह प्रोजेक्ट भी शायद कुछ ऐसा ही है। megakernel में compile करते समय layer boundaries गायब हो जाती हैं, जिससे intermediate result sharing कम होती है और computation बढ़ती है, लेकिन कुल मिलाकर memory round-trip घटने से बड़ा फ़ायदा मिलता है। खासकर convolution networks में शायद कोई sweet spot होगा, लेकिन megakernel वहाँ इसे कैसे handle करता है, यह मुझे नहीं पता.
LLM के लिए अब भी नई-नई उपमाएँ सामने आ रही हैं। कभी-कभी लगता है कि क्या LLM को transistor की तरह सोचा जा सकता है। अभी की स्थिति मानो उस दौर जैसी है जब कमरे जितने बड़े कंप्यूटर punch cards के साथ सिर्फ multiplication करते थे। अगर एक साथ 10 लाख o3-pro queries चलाई जा सकें, तो क्या-क्या हो सकता है, यह कल्पना करना रोचक है.
यह प्रोजेक्ट CMU (Carnegie Mellon) से निकला है। Stanford की Hazy Research ने भी megakernel पर एक ब्लॉग No Bubbles लिखा है। इस क्षेत्र में इतनी सक्रिय प्रतिस्पर्धा देखना प्रभावशाली है। (अतिरिक्त) "mirage" प्रोजेक्ट की बड़ी तस्वीर पर एक paper भी है, लेकिन उसमें megakernel approach शामिल नहीं है paper link
पोस्ट के लेखक ने खुद जवाब दिया। Stanford के साथ शोध समानांतर रूप से आगे बढ़ रहा है, इस बात से सहमति है। मुख्य अंतर यह है कि उनका फ़ोकस automated megakernel generation compiler पर है.
Hazy Research का ThunderKittens भी एक बहुत cool library है। हाल में NVIDIA GPU models का अधिकतम उपयोग करने के लिए formalization, pipelining, divide-and-conquer, efficiency maximization, और dedicated compiler/DSL development पर बहुत प्रयास केंद्रित हुए हैं.
Qwen 8B के performance numbers अगर verify हो जाएँ, तो वे काफ़ी प्रभावशाली होंगे। यह पहले के megakernel तरीकों की तुलना में अधिक practical लगता है। हर SM पर एक kernel बनाए रखने वाला यह तरीका पुराने Larrabee की याद दिलाता है। अगर CUDA की जगह traditional process-thread-SIMD route अपनाया गया होता, तो दुनिया आज कैसी होती, यह सोचना दिलचस्प है.
software-based inference की तुलना में pure ASIC approach से fixed-function LLM बनाने का विचार। क्या इसमें cost advantage होगा? क्या software side पर कुछ अतिरिक्त handling या fine-tuning layers दी जा सकती हैं? चूँकि हम लगभग ‘काफ़ी अच्छा’ स्तर तक पहुँच चुके हैं, इसलिए हो सकता है अगले 2~4 सालों में specialized chips पर models को fix करके उपयोग करने जैसे फ़ैसले लिए जाएँ। सवाल यह है कि ultra-specialized hardware के फ़ायदे आख़िर किस मोड़ पर सचमुच चमकेंगे।