AMD GPU को ‘brrr’ स्पीड से चलाने का तरीका

• HipKittens AMD GPU की छिपी हुई प्रदर्शन क्षमता निकालने के लिए डिज़ाइन किया गया programming primitives का संग्रह है, जो memory access, scheduling और cache reuse को optimize करता है
• AMD MI355X GPU में 256 compute units और 8 chiplets (XCD) की संरचना है, और यह बड़ी register file तथा सूक्ष्म matrix core instructions प्रदान करता है
• NVIDIA के विपरीत AMD में register reallocation, asynchronous matrix instructions, mbarrier नहीं हैं, इसलिए wave specialization की जगह 8-wave ping-pong और 4-wave interleave scheduling अधिक प्रभावी है
• HipKittens chiplet-aware (grid) scheduling के जरिए L2 और LLC cache locality को बेहतर बनाता है, और GEMM व Attention ऑपरेशनों में अधिकतम bandwidth और TFLOPS सुधार हासिल करता है
• यह तरीका AMD GPU ecosystem की software maturity की कमी को पूरा करता है और विविध hardware-आधारित AI computing scalability बढ़ाने की नींव देता है

---

AMD CDNA GPU आर्किटेक्चर और प्रदर्शन विशेषताएँ

• AMD MI355X GPU में 256 compute units (CU) हैं, और हर CU 4 SIMD से बना है
  • एक SIMD 64-thread वाली wave चलाता है, जो NVIDIA के 32-thread warp के मुकाबले है
• MI355X में B200 की तुलना में लगभग 70% SRAM (165KB) है, और इसमें asynchronous matrix multiply instructions, register reallocation, tensor memory acceleration, mbarrier जैसी सुविधाएँ नहीं हैं
• दूसरी तरफ यह 2 गुना बड़ी register file और 60% अधिक processors (256 CU बनाम 160 SM) देता है
  • यह छोटे और सूक्ष्म matrix core instructions को support करता है, और सीधा global→shared memory load (TMA-जैसी) सुविधा भी मौजूद है
• AMD ने 8 chiplets (XCD) वाली chiplet architecture अपनाई है, जहाँ हर XCD की अपनी अलग L2 cache है और उसके ऊपर LLC cache मौजूद है
• तालिका के अनुसार MI355X में BF16 2.5 PFLOPs, MXFP8 5.0 PFLOPs, MXFP6 10.1 PFLOPs की compute performance, और 288GB memory capacity, 8TB/s bandwidth है

AMD के लिए kernel design की चुनौतियाँ

• Memory access optimization: HIPCC compiler की सीमाओं और गैर-सार्वजनिक I/O व्यवहार के कारण data layout और swizzle pattern का design महत्वपूर्ण है
• Processor के अंदर scheduling: AMD में shared memory की बजाय register file और छोटे matrix instructions का उपयोग करना पड़ता है
• Processor के बीच scheduling: chiplet-आधारित संरचना के कारण cache-level NUMA effect को ध्यान में रखकर काम बाँटना जरूरी है

HipKittens के memory access pattern

• HipKittens(HK) tile को मूल data unit की तरह इस्तेमाल करता है, और PyTorch-जैसे operation functions देता है
  • tile को data type, size, layout से परिभाषित किया जाता है, और C++ template metaprogramming के जरिए अलग-अलग inputs को support किया जाता है
• Register scheduling: HIPCC कुछ खास registers को MFMA input के रूप में इस्तेमाल नहीं कर सकता, इसलिए HK explicit register pinning सुविधा देता है
  • इससे developer खुद register तय करके अधिकतम performance वाला kernel लिख सकता है
• Register layout: AMD में data type और matrix shape के अनुसार layout बदलता है, इसलिए एक ही swizzle pattern पर्याप्त नहीं होता
  • उदाहरण के लिए 16×16 bf16 tile और 16×32 bf16 tile को अलग swizzle pattern चाहिए
• Instruction phase structure: AMD की shared memory instructions में discontinuous phase groups होते हैं, और इनका आंतरिक documentation अपर्याप्त है
  • HK इसके लिए reverse-engineered solver प्रदान करता है
• Address generation: AMD asynchronous HBM→shared memory load को support करता है, और HBM address swizzle के जरिए optimization करता है

Processor के अंदर scheduling: Wave pattern

• Wave specialization NVIDIA पर प्रभावी है, लेकिन AMD में register reallocation की कमी के कारण performance गिरती है
  • Producer wave अनावश्यक registers घेर लेता है, और Consumer wave में registers कम पड़ने से spill होता है
• HK के प्रयोगों के अनुसार wave specialization AMD पर arithmetic intensity घटाता है और memory bottleneck पैदा करता है
  • उदाहरण: GEMM में HK 0/8 configuration 1605 TFLOPs देता है, जबकि CUTLASS 1570 TFLOPs देता है
• वैकल्पिक scheduling pattern
  • 8-wave ping-pong: दो wave बारी-बारी से memory/compute cluster चलाते हैं
  • 4-wave interleave: एक wave memory और compute को बारीकी से interleave करता है
  • 8-wave में code सरल रहता है, जबकि 4-wave अधिक सूक्ष्म लेकिन लंबा code बनाता है
  • GEMM और Attention Forward में 8-wave ने SoTA-स्तर का performance हासिल किया

Processor के बीच scheduling: Chiplet-aware तरीका

• AMD MI355X में 8 XCD chiplets हैं, और हर chiplet की अपनी स्वतंत्र L2 cache है
  • thread blocks को round-robin तरीके से chiplets पर बाँटा जाता है, इसलिए grid order सीधे cache reuse efficiency को प्रभावित करता है
• साधारण row-major layout में L2 cache reuse कम रहता है, जिससे bandwidth loss होता है
  • उदाहरण: L2 55%, LLC 95%, 15.1 TB/s, 1113 TFLOPs
• HK ने chiplet-aware (grid) scheduling अपनाई, जिससे L2 और LLC cache locality दोनों का लाभ मिलता है
  • thread blocks को output matrix के आस-पास वाले क्षेत्रों के समूह के रूप में व्यवस्थित कर input data reuse को अधिकतम किया जाता है

वास्तविक kernel उदाहरण

• Attention Forward और BF16 GEMM kernel का hot loop HK की 8-wave ping-pong scheduling का उपयोग करता है
  • हर loop बारी-बारी से Compute–Memory cluster चलाता है और schedule barrier से synchronize होता है
  • code उदाहरण में mma_AtB, load, exp2, col_sum जैसे HK operations बार-बार उपयोग होते हैं

निष्कर्ष: Multi-silicon AI युग में AMD

• HipKittens ने AMD CDNA3·CDNA4 पर प्रतिस्पर्धी performance हासिल की
  • तीन मुख्य बिंदु: optimized memory access, AMD-केंद्रित wave scheduling, chiplet-aware grid scheduling
• HK kernels ने AMD के मानकों पर सर्वोच्च performance हासिल की, और NVIDIA Blackwell kernels के साथ भी प्रतिस्पर्धी स्तर दिखाया
• AI computing में विविधता के लिए AMD GPU की पहुँच बढ़ाना जरूरी है, और HipKittens इसके लिए मुख्य software foundation देता है
• AMD के HIPCC register scheduling में सुधार को भविष्य का महत्वपूर्ण विकास क्षेत्र माना गया है

[GN⁺]

Hacker News राय

• HipKittens पर चर्चा देखना उपयोगी रहेगा
• इसी रिसर्च पर HipKittens: Fast and furious AMD kernels नाम की एक और पोस्ट भी है। उसमें George Hotz और AMD कर्मचारियों की टिप्पणियाँ हैं
• अकादमिक जगत में ऐसे मुद्दों पर काम होना अच्छा है, लेकिन आखिरकार यह AMD के भीतर ही हल किया जाना चाहिए
  • मेरा मानना है कि hardware कंपनियों को hardware पर ही ध्यान देना चाहिए। तभी incentive शुद्ध रहते हैं। performance 20% कम भी हो, तब भी वह बेहतर है
  • पूरी तरह सहमत। AMD ने यह समस्या 10 साल पहले टाल दी थी और अब जाकर इसकी भरपाई करने की कोशिश कर रहा है। hardware शानदार है, लेकिन firmware लिखने की क्षमता की कमी के कारण उसकी पूरी क्षमता सामने नहीं आ पा रही
  • लेकिन जहाँ तक मुझे पता है, इस रिसर्च टीम ने Nvidia GPU के लिए भी इसी तरह का software बनाया था। लगता है बेहतरीन शोधकर्ता अपनी विशेषज्ञता का अच्छा उपयोग कर रहे हैं
  • मेरी जानकारी में AMD पहले से कई स्तरों पर इस समस्या पर काम कर रहा है, और tinycorp के साथ भी सहयोग कर रहा है
• पोस्ट पढ़कर लगता है कि AMD GPU की architectural complexity के कारण optimization कठिन है। लेकिन लंबी अवधि में AMD का approach बेहतर scale कर सकता है। Nvidia जहाँ 2 chiplet इस्तेमाल करता है, वहीं AMD 8 chiplet architecture उपयोग करता है, इसलिए memory locality की समस्या है। भविष्य में chiplet की संख्या और बढ़ सकती है, इसलिए आज की इस complexity को संभालने का अनुभव लंबे समय में फायदेमंद हो सकता है
  • AMD में high performance के लिए warp specialization की ज़रूरत नहीं पड़ती, इसलिए programming अधिक सरल है
• कई developers ने AMD GPU को आम developers के लिए ‘go brrr’ बनाने की कोशिश की, लेकिन असफल रहे। समझ नहीं आता कि AMD खुद software समस्याएँ क्यों हल नहीं करता। अब उनके पास पैसे की भी कमी नहीं है, इसलिए developers को hire न करना कोई बहाना नहीं है। datacenter GPU भी बुरे नहीं हैं, लेकिन व्यक्तिगत ML·AI प्रयोगों के लिए अब भी Nvidia बहुत बेहतर है। मुझे लगता है कि मेरा 5 साल पुराना RTX 3090 अब तक आए किसी भी AMD consumer GPU से बेहतर है
  • AMD का developer experience भयानक है। वे driver crash bug report तक स्वीकार नहीं करते
  • मैंने हाल ही में NVidia 5090 से AMD R9700 32GB के दो कार्डों पर inference server शिफ्ट किया, और अनुभव पूरी तरह सकारात्मक रहा। Fedora kernel पर यह DKMS configuration के बिना तुरंत चल पड़ा, और ROCm के साथ container जोड़ना भी आसान था। सिर्फ Ollama और Storyteller की settings बदलनी पड़ीं। CUDA से कहीं अधिक आरामदायक अनुभव था
  • Nvidia तो Unreal Engine fork तक खुद maintain करता है। AMD उस स्तर पर प्रतिस्पर्धा के आसपास भी नहीं है
  • hardware कंपनियों में Nvidia ही एकमात्र है जो software engineers को competitive compensation देता है। AMD में अब भी software को ‘असली काम’ न मानने वाली संस्कृति बची हुई है, और ऐसी inertia बदलना आसान नहीं है
• Mojo के पास AMD GPU पर developer experience (devX) सुधारने का एक विचार था, उसकी प्रगति क्या हुई यह जानने की उत्सुकता है
• समझ से बाहर है कि AMD software सुधार पर अरबों डॉलर निवेश क्यों नहीं कर रहा। Nvidia दुनिया की सबसे मूल्यवान कंपनी है, और AMD उसका एकमात्र प्रतिद्वंद्वी है
  • AMD कोशिश तो कर रहा है, लेकिन हर साल hardware refresh करने वाली organizational culture को software-केंद्रित culture में बदलना कठिन है। software, hardware की तरह सीधे revenue नहीं लाता, इसलिए management उसकी priority कम रखता है। ऊपर से बाहरी vendors open source के रूप में code दे देते हैं, जो short term में अच्छा लगता है लेकिन long term quality को नुकसान पहुँचाता है। hardware trend एक बार भी छूट जाए तो competitor से पीछे छूटने का जोखिम बहुत बड़ा होता है
  • मैंने कई GPU vendors में काम किया है, और Nvidia ही अकेली कंपनी है जो software को asset मानकर निवेश करती है। बाकी कंपनियाँ उसे सिर्फ लागत समझती हैं
• व्यक्तिगत रूप से मुझे “go brr” meme पसंद नहीं, लेकिन Stanford जैसी जगह पर इसका इस्तेमाल देखना मज़ेदार है
  • दरअसल, 1 साल पहले ThunderKittens announcement के समय भी “go brr” इस्तेमाल किया गया था
  • अगर ऐसे meme विश्वविद्यालय के आधिकारिक चैनलों पर दिखने लगें, तो शायद यह trend खत्म हो जाने का संकेत है
• project शानदार है, लेकिन सवाल यही है कि AMD यह काम खुद क्यों नहीं कर रहा। लगता है AMD अब भी mature software stack के महत्व को ठीक से नहीं समझता। CUDA की तरह ऐसा integrated stack चाहिए जो सभी कार्डों पर काम करे। कभी मुझे विश्वास था कि AMD एक दिन बराबरी कर लेगा, लेकिन अब लगभग उम्मीद छोड़ चुका हूँ
• project अच्छा है, लेकिन पोस्ट खुद अजीब ढंग से लिखी हुई लगती है
  • लेख बहुत असहज लगता है। या तो AI पर ज़रूरत से ज़्यादा निर्भर किया गया है, या AI जैसी शैली की नकल की गई है। “part one देखिए” या “AMD GPU को go brr कैसे बनाएं” जैसी पंक्तियाँ बार-बार दोहराई गई हैं। तकनीकी हिस्सों में भी, जिन बातों को graph से समझाया जा सकता था, उन्हें 100 लाइन के code में समझाने की कोशिश खास तौर पर खटकती है