C++ और CUDA का उपयोग करके शुरुआत से LLM inference engine बनाना

 (andrewkchan.dev)
15 पॉइंट द्वारा GN⁺ 2024-12-16 | 1 टिप्पणियां | WhatsApp पर शेयर करें
  • C++ और CUDA का उपयोग करके बिना किसी लाइब्रेरी के LLM inference engine बनाने का तरीका
  • इसके जरिए LLM inference की पूरी स्टैक को समझा जा सकता है, और यह महसूस किया जा सकता है कि अलग-अलग optimization inference speed को कैसे प्रभावित करते हैं
  • लक्ष्य: एक single CPU + GPU server पर single batch में तेज inference करने वाला मॉडल लागू करना और llama.cpp से तेज token processing speed हासिल करना

1. LLM आर्किटेक्चर और inference का अवलोकन

  • अधिकांश प्रमुख LLM लगातार transformer blocks का उपयोग करने वाली एक ही आर्किटेक्चर का पालन करते हैं।
  • मॉडल लोडिंग में एक customizable transformer block class को परिभाषित करना, उसे sequence के रूप में व्यवस्थित करना, और safetensors weights से initialize करना शामिल है।
  • inference मुख्य रूप से single batch में होता है, और "decode stage" execution का अधिकांश हिस्सा लेता है।

1.1 inference का अवलोकन

  • inference को prefill stage और decode stage में बाँटा जाता है: prefill stage में दिए गए prompt tokens को मॉडल में भेजकर KV cache भरी जाती है, और decode stage में मॉडल को बार-बार चलाकर tokens जनरेट किए जाते हैं
    • Prefill stage: prompt tokens को process करते हुए KV cache initialize करना
    • Decode stage: एक बार में एक token generate करना
  • KV cache: पिछले key/value pairs को store करके पिछले context के साथ attention को तेज़ी से compute करता है
  • मॉडल का forward pass embedding table का उपयोग करके token IDs को embedding vectors में map करता है, और transformer block sequence के जरिए state को transform करता है

1.2 bottleneck और benchmark

  • bottleneck: आधुनिक hardware में memory bandwidth limiting factor है
    • मॉडल inference में हर token generate करने के लिए पूरे मॉडल को पढ़ना पड़ता है, इसलिए computation की तुलना में memory bandwidth बड़ी बाधा बनती है
  • मॉडल quantization inference speed सुधारने में प्रभावी है
  • सैद्धांतिक अधिकतम token throughput hardware के अनुसार अलग होता है, और वास्तविक performance को कई inference engines के जरिए जाँचा जा सकता है
  • सैद्धांतिक speed limit:
    • AMD EPYC 7702P: अधिकतम 13.6 tok/s (FP16 आधार)
    • RTX 4090: अधिकतम 67.1 tok/s (FP16 आधार)
  • benchmark:
    • llama.cpp: CPU 8.7 tok/s, GPU 61 tok/s
    • calm: GPU 66 tok/s

2. CPU-आधारित inference

  • CPU पर शुरुआती implementation single thread में किया गया है, और केवल FP32 weights को support करता है
  • multithreading के जरिए code parallelization शुरू किया जा सकता है, और SIMD का उपयोग करके performance बेहतर की जा सकती है

2.1 multithreading

  • OpenMP का उपयोग करके matrix-vector multiplication (matmul) और multi-head attention को parallelize कर performance बेहतर की जाती है
  • optimization परिणाम: speed 0.6 tok/s → 4.4 tok/s तक सुधरी

2.2 weight quantization और SIMD optimization

  • quantization: FP32 weights को FP16 में quantize करके memory usage आधा किया जाता है और performance सुधरती है
  • SIMD: AVX2 का उपयोग करके 8 FP32 values को एक साथ process करने के लिए optimize किया गया
  • परिणाम: 8.4 tok/s हासिल

3. GPU-आधारित inference

  • मॉडल को FP16 में quantize करके RTX 4090 पर लोड किया जा सकता है और GPU inference implementation शुरू की जा सकती है
  • CUDA का उपयोग करके C++ functions (kernels) को GPU पर parallel में चलाया जा सकता है

3.1 CUDA में सरल porting

  • CPU operations को 1-1 CUDA kernels में बदलकर GPU backend लागू किया जा सकता है
  • CUDA kernels asynchronous रूप से चलते हैं, लेकिन एक ही stream में क्रमवार execute होते हैं
  • समस्या: thread inefficiency के कारण GPU resources का पर्याप्त उपयोग नहीं हो पाता → 2.9 tok/s के साथ धीमा

3.2 बेहतर matrix multiplication (matmul)

  • matrix multiplication CPU पर runtime का बड़ा हिस्सा लेती है, और OpenMP के जरिए optimize की जा सकती है
  • GPU में प्रति block 1 row process करके thread utilization बढ़ाया जा सकता है
  • optimization तरीके:
    1. एक block एक row process करे, और block के अंदर threads मिलकर computation करें
    2. warp reduction लागू करना
  • परिणाम: speed बढ़कर 51.7 tok/s

3.3 kernel fusion और अतिरिक्त optimization

  • kernels को fuse करके performance बढ़ाई जा सकती है
    • kernel fusion: लगातार होने वाले operations को एक kernel में मिलाकर memory access और computation time को कम करना
  • memory access pattern optimization और space reuse के जरिए 56.1 tok/s हासिल

3.4 Attention optimization और लंबे context की प्रोसेसिंग

  • समस्या: लंबे context में attention kernel performance bottleneck बनता है
  • समाधान:
    1. memory access optimization: लगातार memory blocks पढ़ने के लिए redesign
    2. atomicAdd की जगह shared memory का उपयोग करके missing floating-point values की समस्या हल करना
  • optimization परिणाम:
    • short context: 63.8 tok/s (llama.cpp के 61.0 tok/s से तेज)
    • long context: 58.8 tok/s हासिल

3.5 KV cache quantization और compiler optimization समस्याएँ

  • KV cache को FP16 में quantize करने पर performance गिरावट होती है (compiler optimization की कमी के कारण)
  • समाधान: manually loop unrolling और memory prefetching लागू करना
  • परिणाम: FP32 की तुलना में लगभग 2x speed improvement और long context performance 58.8 tok/s पर बरकरार

4. आगे के सुधार की दिशा

  • prompt prefill optimization: कई tokens को एक साथ process करके first-token generation time कम करना
  • Attention kernel fusion: FlashAttention जैसी optimization techniques लागू करना
  • उच्चतर quantization: FP8, INT8, INT4 लागू करना और activation/cache quantization
  • kernel optimization: memory bandwidth और computation efficiency को अधिकतम करने वाली advanced techniques अपनाना
  • लाइब्रेरी का उपयोग: cuDNN, cuBLAS जैसी libraries का उपयोग करके optimization time घटाना

परिणाम सारांश:

  • CPU और GPU पर विभिन्न optimizations के जरिए 63.8 tok/s speed हासिल
  • llama.cpp और calm के बराबर या उससे बेहतर performance दर्ज
  • बिना लाइब्रेरी के केवल C++ और CUDA से high-performance LLM inference engine लागू किया गया

संबंधित पढ़ाई

1 टिप्पणियां

 
GN⁺ 2024-12-16
Hacker News प्रतिक्रियाएँ
  • लेखक अपने ब्लॉग पोस्ट पर ध्यान जाने से खुश है और feedback सुनना चाहता है
  • एक पाठक ने लेख की बहुत प्रशंसा की और पूछा कि इसे लिखने में कितना समय लगा
    • GPGPU क्षेत्र में काम करने वाले व्यक्ति के रूप में वह भी ऐसा ही लेख लिखना चाहता है, लेकिन लगने वाले समय की अनिश्चितता की वजह से हिचक रहा है
  • एक अन्य पाठक का मानना है कि कोड tensor cores या wgmma निर्देशों का उपयोग नहीं करता
    • उसका कहना है कि इस तरह की programming मुश्किल होती है क्योंकि इसमें कई काम एक साथ संभालने पड़ते हैं
    • वह यह भी कहता है कि bandwidth limits की वजह से शायद अतिरिक्त computation की ज़रूरत न पड़े
    • उसका आकलन है कि ब्लॉग का कोड दूसरे accelerators पर port करते समय बेहतर काम करेगा
    • उसे चिंता है कि wgmma का उपयोग करने से Nvidia की पीढ़ियों के बीच portability कम हो सकती है
  • एक और पाठक इसी तरह की Python सामग्री ढूंढ रहा है और उसे अपनी टीम के साथ साझा करना चाहता है
    • वह performance से अधिक conceptually complete और tutorial शैली में संक्षिप्त सामग्री चाहता है
  • एक उपयोगकर्ता अपने Mistral version और tokens/second performance की तुलना करना चाहता है
    • README के quantization section को देखने की सलाह दी गई
  • यह राय भी है कि __shfl_down आजकल warp synchronization समस्याओं के कारण अनुशंसित नहीं है