DeepGEMM: सूक्ष्म स्केलिंग के साथ स्वच्छ और कुशल FP8 GEMM कर्नेल

DeepGEMM

DeepGEMM, FP8 general matrix multiplication (GEMM) के लिए एक लाइब्रेरी है, जो DeepSeek-V3 में प्रस्तावित सूक्ष्म स्केलिंग का समर्थन करती है। यह लाइब्रेरी सामान्य और Mix-of-Experts (MoE) grouped GEMM, दोनों का समर्थन करती है, और CUDA में लिखी गई है, इसलिए इंस्टॉल करते समय compilation की आवश्यकता नहीं होती। यह NVIDIA Hopper tensor core का समर्थन करती है, और FP8 tensor core accumulation की अशुद्धि को हल करने के लिए CUDA core पर 2-स्तरीय accumulation का उपयोग करती है। यह CUTLASS और CuTe की कुछ अवधारणाओं का उपयोग करती है, लेकिन templates या algebra पर निर्भरता को न्यूनतम रखकर सरलता बनाए रखती है। लगभग 300 लाइनों के कोड से बनी एक मुख्य kernel function के साथ, यह Hopper FP8 matrix multiplication और optimization तकनीकों को सीखने के लिए एक उपयुक्त संसाधन है। हल्के डिज़ाइन के बावजूद, यह विभिन्न matrix shapes में expert-tuned लाइब्रेरी के बराबर या उससे बेहतर प्रदर्शन देती है.

प्रदर्शन

DeepSeek-V3/R1 inference में उपयोग होने वाले सभी shapes को H800 SXM5 पर NVCC 12.8 के साथ टेस्ट किया गया। सभी speedup metrics की गणना CUTLASS 3.6 पर आधारित internally optimized implementation की तुलना में की गई। कुछ shapes में प्रदर्शन अच्छा नहीं हो सकता है, और optimization PR का स्वागत है.

सामान्य GEMM (dense model)

• विभिन्न matrix sizes में DeepGEMM के प्रदर्शन को मापने पर, कुछ खास sizes में अधिकतम 2.7x speedup देखा गया.

MoE मॉडल के लिए grouped GEMM (continuous layout)

• groups की संख्या और प्रत्येक group के matrix size के आधार पर अधिकतम 1.2x speedup देखा गया.

MoE मॉडल के लिए grouped GEMM (mask layout)

• mask layout का उपयोग करके अधिकतम 1.2x speedup देखा गया.

त्वरित शुरुआत

आवश्यकताएँ

• Hopper architecture GPU, sm_90a समर्थन आवश्यक
• Python 3.8 या उससे ऊपर
• CUDA 12.3 या उससे ऊपर (सर्वश्रेष्ठ प्रदर्शन के लिए 12.8 या उससे ऊपर अनुशंसित)
• PyTorch 2.1 या उससे ऊपर
• CUTLASS 3.6 या उससे ऊपर

डेवलपमेंट

• submodule clone करना, symbolic link बनाना, JIT compilation, और सभी GEMM implementations की testing सहित डेवलपमेंट प्रक्रिया का विवरण.

इंस्टॉलेशन

• deep_gemm को Python project में import करके उपयोग किया जा सकता है.

इंटरफ़ेस

ध्यान दें

• यह लाइब्रेरी केवल GEMM kernels शामिल करती है और केवल NT format को समर्थन देती है। transpose या अन्य FP8 casting कार्यों को अलग से implement करना होगा.

सामान्य dense GEMM (non-grouped)

• मूल non-grouped FP8 GEMM चलाने के लिए functions प्रदान करती है.

grouped GEMM (continuous layout)

• MoE मॉडल में उन scenarios के लिए डिज़ाइन किया गया है जहाँ experts एक ही shape साझा करते हैं.

grouped GEMM (mask layout)

• inference decoding चरण में mask tensor देकर केवल मान्य भागों की गणना की जाती है.

यूटिलिटी

• विभिन्न utility functions और environment variables प्रदान करती है, जो performance optimization में मदद करते हैं.

अनुकूलन

persistent warp specialization

• CUTLASS डिज़ाइन का पालन करते हुए, data movement, tensor core MMA instructions, और CUDA core promotion को overlap किया जाता है.

Hopper TMA सुविधा

• data movement को तेज़ करने के लिए TMA का उपयोग किया जाता है.

सामान्य सूक्ष्म अनुकूलन

• विभिन्न optimization तकनीकों के माध्यम से प्रदर्शन सुधारा जाता है.

एकीकृत और अनुकूलित block scheduler

• सभी non-grouped और grouped kernels के लिए scheduler प्रदान किया जाता है.

पूर्ण JIT डिज़ाइन

• install समय compilation की आवश्यकता न होने वाला JIT डिज़ाइन प्रदर्शन को बेहतर बनाता है.

unaligned block size

• कुछ विशेष shapes में SM utilization को अधिकतम करने के लिए unaligned block size का समर्थन किया जाता है.

FFMA SASS interleaving

• प्रदर्शन सुधारने के लिए FFMA instructions को समायोजित कर warp-level parallelism को बढ़ाया जाता है.

आभार

• DeepGEMM, CUTLASS project से प्रेरित है, और इसके डेवलपर्स के प्रति आभार और सम्मान व्यक्त करता है.

लाइसेंस

• MIT लाइसेंस के तहत जारी किया गया है.