Perplexity-ൽ CuTeDSL

Perplexity-യിൽ, എംബെഡിംഗുകളിൽ നിന്ന് ട്രില്ല്യൺ-പാരാമീറ്റർ LLMകൾ വരെ ഉള്ള മോഡലുകൾ സർവ് ചെയ്യാൻ ഞങ്ങൾ ആഭ്യന്തരമായ Runtime-Optimized Serving Engine (ROSE) ഉപയോഗിക്കുന്നു. Sonar, Search, Embeddings എന്നിവ ഉൾപ്പെടെ എല്ലാ APIകളുടെയും പിന്നിൽ നിലകൊള്ളുന്ന ഞങ്ങളുടെ റൺടൈം, ഗവേഷകർക്ക് ഏറ്റവും അനുയോജ്യമായ മോഡലുകളും ആർക്കിടെക്ചറുകളും അന്വേഷിക്കാൻ ശക്തി നൽകുന്നതിനും, അതേസമയം മികച്ച സ്ഥാനാർഥികളെ വേഗത്തിൽ പ്രൊഡക്ഷനിലേക്ക് വിന്യസിക്കുന്നതിനുമായി നിർമിച്ചതാണ്. ഇത് കൈവരിക്കാൻ, അത്യാധുനിക ആർക്കിടെക്ചറുകൾക്കായുള്ള inference വേഗത്തിൽ നടപ്പാക്കാനും, വിന്യസിക്കാനും, ഓപ്റ്റിമൈസ് ചെയ്യാനും ഞങ്ങൾക്ക് കഴിയേണ്ടതുണ്ട്. ഈ ബ്ലോഗ് പോസ്റ്റിൽ, ഞങ്ങളുടെ inference engine-ലേക്ക് ഏകീകരിച്ചിരിക്കുന്ന CuTeDSL kernels NVIDIA GPUകളിൽ മോഡലുകളെ പരമാവധി പ്രകടനത്തിലേക്ക് എത്തിക്കാൻ ആവശ്യമായ GPU kernels നിർമ്മിക്കാൻ എങ്ങനെ സഹായിക്കുന്നുവെന്ന് നാം പരിശോധിക്കുന്നു.

Perplexity-ലെ വിവിധ ടീമുകൾ നിർമ്മിക്കുന്ന കസ്റ്റം മോഡലുകൾ, AI ഇൻഫറൻസ് ടീം നിർമ്മിച്ച ROSE ഇൻഫറൻസ് എൻജിൻ ഉപയോഗിച്ച് NVIDIA Hopper, Blackwell GPU-കളിൽ ഇൻ-ഹൗസ് ആയി ഹോസ്റ്റ് ചെയ്യപ്പെടുന്നു. ഭാഷാ ഡീകോഡിംഗിനും ക്ലാസിഫിക്കേഷനും ഒരേസമയം കസ്റ്റമൈസ് ചെയ്ത Llama മോഡലുകൾ മെച്ചപ്പെട്ട രീതിയിൽ നൽകുന്നതിനാണ് ROSE ആദ്യമായി നിർമ്മിച്ചത്, NVIDIA Triton ഇൻഫറൻസ് സെർവറുമായി പൊരുത്തപ്പെടുന്ന ഒരു ഇന്റർഫേസും ഇതിലുണ്ടായിരുന്നു. കാലക്രമേണ, Perplexity-യുടെ ഉൽപ്പന്നങ്ങൾക്ക് കരുത്ത് പകരുന്നതിനായി വിപുലമായ LLM-കൾക്കും റാങ്കിംഗ്, ക്ലാസിഫിക്കേഷൻ, സ്കോറിംഗ്, മറ്റ് ട്രാൻസ്ഫോർമർ അധിഷ്ഠിത മോഡലുകൾക്കും സേവനം നൽകുന്ന പൂർണ്ണമായ ഒരു ഇൻഫറൻസ് എൻജിനായി ഇത് വളർന്നു. ലളിതമായിരിക്കുക എന്നതാണ് ROSE-ന്റെ ഡിസൈൻ ലക്ഷ്യം: കാലക്രമേണ, മികച്ച പ്രകടനം കൈവരിക്കുന്നതിനായി പരിശ്രമങ്ങൾ നിക്ഷേപിക്കപ്പെടുന്നു, എന്നാൽ ഞങ്ങളുടെ മോഡലുകൾക്കായുള്ള എക്കാലത്തെയും വികസിച്ചുകൊണ്ടിരിക്കുന്ന ആർക്കിടെക്ചറുകളോടും യൂസ് കേസുകളോടും പൊരുത്തപ്പെടുക എന്നതാണ് പ്രധാന മുൻഗണന.

ഡാറ്റാ തരം, ഹെഡ് ഡൈമൻഷനുകൾ, കമ്പൈൽ-ടൈം കോൺസ്റ്റന്റുകൾ, ലോഞ്ച് സെറ്റിംഗുകൾ എന്നിവയാകെ കേർണൽ കോൺഫിഗറേഷൻ സ്പേസ്.

ഉയർന്ന തലത്തിൽ, ക്ലയന്റുകൾക്ക് സംവദിക്കാൻ കഴിയുന്ന ഒരു ഇന്റർഫേസിലേക്ക് ഒരു മോഡലിനെ അനുയോജ്യമാക്കുന്ന ഒരു engine ROSE-ൽ അടങ്ങിയിരിക്കുന്നു. എൻജിനുകൾ വളരെ എളുപ്പത്തിൽ വീണ്ടും ഉപയോഗിക്കാവുന്നവയാണ്, അവ അഭ്യർത്ഥന ഷെഡ്യൂളിംഗും ബാച്ചിംഗും കൈകാര്യം ചെയ്യുകയും അടിസ്ഥാന മോഡലുകളിൽ നിന്ന് ഡീകോഡ് ചെയ്ത ടോക്കണുകൾ, എംബെഡിംഗുകൾ അല്ലെങ്കിൽ സ്കോറുകൾ വീണ്ടെടുക്കുകയും ചെയ്യുന്നു. ഉപകരണങ്ങളുടെ പ്രാരംഭ സജ്ജീകരണവും, പ്രോസസ്സറുകൾക്കും നോഡുകൾക്കുമിടയിലുള്ള ആശയവിനിമയ ചാനലുകളും, വെയിറ്റ് ലോഡിംഗും ഇവയെല്ലാം അവയാണ് കൈകാര്യം ചെയ്യുന്നത്. LLM-കൾക്കായി, ബാച്ചിംഗ്, ചങ്കിംഗ്, സാമ്പിളിംഗ് എന്നിവയ്ക്ക് എൻജിൻ ഉത്തരവാദിയാണ്. മുഴുവൻ, ലീനിയർ അറ്റൻഷൻ എന്നിവയ്ക്കായുള്ള KV സ്റ്റോറേജ് അലോക്കേഷനും, പ്രീഫിക്സ് മാച്ചിംഗിനും ഇത് ഉത്തരവാദിയാണ്. എംബെഡിംഗ് എൻജിനുകൾ ലളിതമാണ്, കാരണം അവ ഓൺലൈൻ ബാച്ചിംഗ് മാത്രമേ നടത്തുന്നുള്ളൂ. ആക്സിലറേറ്ററിൽ മോഡൽ സമാന്തരമായി പ്രവർത്തിക്കുമ്പോൾ അടുത്ത ഘട്ടത്തിലുള്ള ജോലി എൻജിൻ തയ്യാറാക്കുന്ന രീതിയിൽ എക്സിക്യൂഷൻ ഒപ്റ്റിമൈസ് ചെയ്തിരിക്കുന്നു.

CuTeDSL, Triton, CUDA, CUTLASS, cuBLAS എന്നിവ ഉപയോഗിച്ച് നിർമ്മിച്ച GPU കർണലുകളെ റാപ്പ് ചെയ്യുന്ന, മോഡലിംഗിനായുള്ള കസ്റ്റം ലെയറുകളുടെ സമ്പന്നമായ ഒരു ശേഖരം ROSE നൽകുന്നു. മാട്രിക്സ് ഗുണനത്തിനും അറ്റൻഷൻ ഇംപ്ലിമെന്റേഷനുകൾക്കുമായി, ഞങ്ങൾ പ്രധാനമായും CUTLASS-ൽ നിന്നും cuBLAS-ൽ നിന്നുമുള്ള NVIDIA കർണലുകളെയാണ് ആശ്രയിക്കുന്നത്. എംബെഡിംഗ്, നോം, MoE, ആക്ടിവേഷൻ കർണലുകൾ ഉൾപ്പെടെയുള്ള പൈപ്പ്‌ലൈനിന്റെ ബാക്കി ഭാഗങ്ങൾ ROSE-ന്റെ കർണൽ ലൈബ്രറിയുടെ ഭാഗമാണ്. ഞങ്ങൾ ഇൻഫറൻസ് നടത്തുന്ന ഹാർഡ്‌വെയറിനെ മികച്ച രീതിയിൽ പിന്തുണയ്ക്കുന്നതിന്, MoE റൂട്ടിംഗ്, ഡിസ്പാച്ച്, കമ്പൈൻ എന്നിവയ്ക്കായി സ്വന്തമായി ഡിവൈസ്-ടു-ഡിവൈസ് കമ്മ്യൂണിക്കേഷൻ പ്രിമിറ്റീവുകൾ ROSE-ൽ സജ്ജീകരിച്ചിരിക്കുന്നു. മറ്റ് കർണലുകൾ താരതമ്യേന ലളിതമാണെങ്കിലും, അവ പ്രവർത്തിക്കുന്ന വിവിധ ഹിഡൻ ഡയമെൻഷനുകളിലും ഉപകരണങ്ങളിലും ഉയർന്ന പ്രകടനം കൈവരിക്കുന്നതിനായി അവ പൂർണ്ണമായും പ്രത്യേകവൽക്കരിക്കപ്പെടണം. കൂടാതെ, വിവിധ മോഡലുകളുടെ എല്ലാ നോബുകളും ട്വീക്കുകളും ഉൾക്കൊള്ളുന്നതിനായി ഞങ്ങൾ കംപൈൽ-ടൈം സ്പെഷ്യലൈസേഷനെ ആശ്രയിക്കുന്നു (RMS നോർമിനുള്ള വെയിറ്റ് ബയാസ് vs വെയിറ്റ് ബയാസ് ഇല്ലായ്മ, ആക്ടിവേഷൻ ഫംഗ്ഷനുകൾ മുതലായവ). തൽഫലമായി, അവ നടപ്പിലാക്കാൻ DSL-കൾ ഉപയോഗിക്കാൻ ഞങ്ങൾ തീരുമാനിച്ചു: ഞങ്ങൾ ആദ്യം Triton ഉപയോഗിച്ചാണ് ROSE നിർമ്മിച്ചത്. എന്നിരുന്നാലും, അതിനുശേഷം ഞങ്ങൾ GPU പ്രോഗ്രാമിംഗിനായുള്ള പ്രാഥമിക പരിസ്ഥിതിയായി CuTeDSL സ്വീകരിച്ചു.

ഗണിതശാസ്ത്രപരമായ കാഴ്ചപ്പാടിൽ, ഇൻഫറൻസ് കർണലുകൾ സാധാരണയായി സങ്കീർണ്ണമല്ലാത്തവയാണ്. തൽഫലമായി, ഒരു കർണൽ വികസിപ്പിച്ചെടുക്കുമ്പോൾ, സങ്കീർണ്ണത സാധാരണയായി എന്ത് കമ്പ്യൂട്ടേഷനാണ് ചെയ്യേണ്ടത് എന്ന് പ്രകടിപ്പിക്കുന്നതിലല്ല, മറിച്ച് ഹാർഡ്‌വെയർ അത് കൃത്യമായി എങ്ങനെ ചെയ്യണം എന്നതിലാണ്. മികച്ച കർണലുകൾ നിർമ്മിക്കുന്നതിന്, അടിസ്ഥാന ഹാർഡ്‌വെയറിന് മുകളിൽ ശരിയായ അമർബ്സ്ട്രാക്ഷനുകൾ നൽകുന്ന ഒരു ഭാഷ ഞങ്ങൾക്ക് ആവശ്യമാണ്. കമ്പ്യൂട്ടേഷൻ നടപ്പിലാക്കുന്നതിനുള്ള ഉയർന്ന തലത്തിലുള്ള പ്രവർത്തനങ്ങൾ വ്യക്തമാക്കാൻ ശരിയായ ഭാഷ നമ്മെ പ്രാപ്തരാക്കുന്നു, കൂടാതെ വിപുലമായ കോൺഫിഗറേഷൻ സ്പേസിലെ എല്ലാ പോയിന്റുകളിലും അവയെ ആക്രമണാത്മകമായി സ്പെഷ്യലൈസ് ചെയ്യുകയും ഒപ്റ്റിമൈസ് ചെയ്യുകയും ചെയ്യുന്നു.

CuTe ലേഔട്ട് ആൾജിബ്രയ്ക്കും MLIR-നും മുകളിൽ നിർമ്മിച്ച NVIDIA GPU-കൾക്കായുള്ള Python-അധിഷ്ഠിത DSL ആണ് CuTeDSL. കുറഞ്ഞ തലത്തിലുള്ള ഹാർഡ്‌വെയർ പ്രിമിറ്റീവുകളുടെ നിയന്ത്രണം നിലനിർത്തിക്കൊണ്ട് ഇത് ഉയർന്ന ഒപ്റ്റിമൈസ് ചെയ്ത PTX-ലേക്ക് ജസ്റ്റ്-ഇൻ-ടൈം കംപൈൽ ചെയ്യുന്നു. വികസന വേഗതയെ തടസ്സപ്പെടുത്തുന്ന ചെലവേറിയ കംപൈലേഷൻ സമയങ്ങളില്ലാതെ, CUTLASS C++-ന്റെ അതേ സവിശേഷതകളും പ്രകടനവും വാഗ്ദാനം ചെയ്യാൻ ഇത് ലക്ഷ്യമിടുന്നു. ഞങ്ങൾ സർവേ ചെയ്തതോ പ്രവർത്തിച്ചതോ ആയ ഭാഷകളിൽ നിന്ന് (CUDA, cuTile, TVM IR, Triton എന്നിവ ഉൾപ്പെടെ), ഞങ്ങളുടെ കർണലുകളുടെ വികസനത്തിന് ഏറ്റവും അനുയോജ്യമായത് ഇതാണെന്ന് ഞങ്ങൾ കരുതുന്നു.

മൊത്തത്തിൽ, Triton അല്ലെങ്കിൽ cuTile പോലെ ഉയർന്ന തലത്തിലുള്ളതും സംക്ഷിപ്തവുമായ ഒന്നല്ലെങ്കിലും, CuTe ലേഔട്ട് ആൾജിബ്രയും അതിന്റെ കോമ്പോസിഷനുകളും (cute.Tensor വഴി ലഭ്യമാക്കിയിരിക്കുന്നു) ന്യായമായ പകരക്കാരനാണെന്ന് ഞങ്ങൾ കണ്ടെത്തി. കൂടാതെ, CuTeDSL ഹാർഡ്‌വെയർ പ്രിമിറ്റീവുകളിലേക്ക് പരിധിയില്ലാത്ത ആക്സസ് വാഗ്ദാനം ചെയ്യുന്നു, എല്ലാ എഡ്ജ് കേസുകളിലും ഉയർന്ന പ്രകടനം കൈവരിക്കാൻ ഞങ്ങൾക്ക് ആവശ്യമായ തലത്തിലുള്ള നിയന്ത്രണം ഇത് നൽകുന്നു. CUDA ഉപയോഗിച്ച്, കംപൈൽ-ടൈം സ്പെഷ്യലൈസേഷനുള്ള അവസരങ്ങൾ പരിമിതമായിരുന്നു, കാരണം ഓരോ അധിക പാരാമീറ്ററും കംപൈലേഷൻ സമയങ്ങൾക്ക് അനാവശ്യമായ ചെലവോ അല്ലെങ്കിൽ ബിൽഡ് സിസ്റ്റത്തിന് സങ്കീർണ്ണതയോ ചേർത്തു. ഇതിനു വിപരീതമായി, ഒരു ജസ്റ്റ്-ഇൻ-ടൈം കംപൈൽഡ് DSL ഉപയോഗിച്ച്, ആവർത്തന വേഗതയെ ബലിയർപ്പിക്കാതെ നമുക്ക് കൂടുതൽ ആക്രമണാത്മകമായി സ്പെഷ്യലൈസ് ചെയ്യാം. CUDA ഒരു ഇടനില ഘട്ടമായി ഉപയോഗിച്ച TVM IR-ൽ നിന്ന് വ്യത്യസ്തമായി, PTX-ലേക്ക് താഴ്ത്തുന്ന MLIR-അധിഷ്ഠിത പൈപ്പ്‌ലൈൻ കൂടുതൽ കരുത്തുറ്റതാണെന്നും CUDA ഇടനില ഘട്ടമായി പുറന്തള്ളുന്നതിനേക്കാൾ ഡീബഗ് ചെയ്യാൻ എളുപ്പമാണെന്നും ഞങ്ങൾ കണ്ടെത്തി. കൂടാതെ, PTX നിർദ്ദേശങ്ങൾ റാപ്പ് ചെയ്യാൻ കസ്റ്റം ഓപ്പുകൾ ഉപയോഗിക്കുന്നതിലൂടെ, അടിസ്ഥാന ഹാർഡ്‌വെയറിന്റെ എല്ലാ സവിശേഷതകളും പൂർണ്ണമായി പ്രയോജനപ്പെടുത്താൻ ഞങ്ങൾക്ക് കഴിയുമെന്ന് ഞങ്ങൾ കണ്ടെത്തി.

CUDA കർണലുകളുടെ കംപൈലേഷൻ ചെലവിന്റെ ഭൂരിഭാഗവും സാധാരണയായി കർണലുകളുടെ ടെംപ്ലേറ്റ് എക്സ്പാൻഷൻ ഘട്ടത്തിലാണ് ചെലവഴിക്കുന്നത്, കാരണം PTX ഘട്ടം മുതൽ ഒപ്റ്റിമൈസേഷനും കംപൈലേഷനും നിമിഷങ്ങൾക്കുള്ളിൽ പൂർത്തിയാക്കാൻ കഴിയും. നിരവധി കംപൈൽ-ടൈം പാരാമീറ്ററുകൾ ഉള്ളതിനാൽ, ലളിതമല്ലാത്ത ഒരു കർണലിനായി കംപൈലേഷന്റെ ഈ ഘട്ടം പല മിനിറ്റുകളായി വർദ്ധിച്ചേക്കാം. C++ ടെംപ്ലേറ്റ് കൈകാര്യം ചെയ്യൽ ക്രമമായി നടക്കുമ്പോൾ, IR, PTX തലങ്ങളിൽ മാത്രമേ സമാന്തര കംപൈലേഷനുള്ള അവസരങ്ങൾ പിന്നീട് ഉണ്ടാകൂ എന്ന വസ്തുതയാൽ ഈ തടസ്സം കൂടുതൽ പ്രകടമാകുന്നു. ടെംപ്ലേറ്റ് ഇൻസ്റ്റാൻഷിയേഷനുകൾ പ്രത്യേക കംപൈലേഷൻ യൂണിറ്റുകളിലേക്ക് മാറ്റിക്കൊണ്ട് ഇത് പരിഹരിക്കാൻ വഴികൾ ഉണ്ടെങ്കിലും, അവയ്ക്ക് ബിൽഡ് സിസ്റ്റങ്ങളിൽ കാര്യമായ മാറ്റങ്ങളും ഭാരമേറിയ മാക്രോകളുടെ ഉപയോഗവും ആവശ്യമാണ്. ഇതിനു വിപരീതമായി, ജസ്റ്റ്-ഇൻ-ടൈം (JIT) കംപൈൽ ചെയ്യപ്പെടുന്ന CuTeDSL കർണലുകൾക്ക് സങ്കീർണ്ണമായ ബിൽഡ് നിയമങ്ങൾ ആവശ്യമില്ല കൂടാതെ കംപൈലേഷൻ സമയങ്ങളെ തടസ്സപ്പെടുത്തുന്നില്ല, കാരണം വികസന ചക്രം കോൺഫിഗറേഷൻ സ്പേസിലെ ഒരൊറ്റ പോയിന്റിൽ ശ്രദ്ധ കേന്ദ്രീകരിക്കാൻ കഴിയും, അതേസമയം പിന്നീട് ശ്രദ്ധയോടെയുള്ള യൂണിറ്റ് ടെസ്റ്റുകൾക്ക് കർണലിന്റെ മുഴുവൻ ഡിസൈൻ സ്പേസും കവർ ചെയ്യാൻ കഴിയും.

JIT compilation വികസനസമയത്തിലെ ചെലവുകൾ runtime-ലേക്ക് മാറ്റുന്നുണ്ടെങ്കിലും, അതിന്റെ ഓവർഹെഡ് കൂടുതൽ അംഗീകരിക്കാവുന്നതാണ്. ഒരു മോഡൽ വിന്യസിക്കപ്പെടുമ്പോൾ, അത് മുഴുവൻ configuration space-ലുള്ള കുറച്ച് points മാത്രമാണ് സജീവമാക്കുന്നത്, inference engine-ന്റെ startup/warmup സമയത്ത് second-scale ഓവർഹെഡ് ഉണ്ടാക്കുന്നു. ഇതിന് വിരുദ്ധമായി, ഒരു kernel-ന്റെ എല്ലാ configurations-ഉം pre-compile ചെയ്യുന്നത് നിരവധി മിനിറ്റുകൾ എടുക്കാം. കൂടാതെ, JIT-compiled artifacts cache ചെയ്യുന്നതിലൂടെ ഓവർഹെഡുകൾ കൂടുതൽ കുറയ്ക്കാനും കഴിയും.

CuTe, CUTLASS എന്നിവയുടെ കാര്യത്തിലെന്നപോലെ, ടെംപ്ലേറ്റ് അധിഷ്ഠിത ഡിസ്പാച്ചും സ്പെഷ്യലൈസേഷനും ഉപയോഗിക്കുമ്പോൾ പിഴവുകളുടെ ഉത്ഭവം കണ്ടെത്താനുള്ള ബുദ്ധിമുട്ട് കാരണം CUDA-യുടെ ഭാരമേറിയ കംപൈലേഷൻ ചെലവ് കൂടുതൽ പ്രകടമാകുന്നു. ഒരു DSL ഉപയോഗിച്ച്, പിഴവുകൾ അത് ഉണ്ടാക്കിയ കോഡ് വരിയിലേക്ക് ചൂണ്ടിക്കാണിക്കാൻ കഴിയും, ഇത് ഡീബഗ്ഗിംഗ് വളരെ എളുപ്പമാക്കുന്നു. കൂടാതെ, ഇടനില MLIR ഘട്ടങ്ങൾ വഴിയോ അല്ലെങ്കിൽ ജനറേറ്റഡ് PTX-ൽ ഘടിപ്പിച്ചിട്ടുള്ള ഡീബഗ്ഗിംഗ് വിവരങ്ങളെ ആശ്രയിച്ചോ ജനറേറ്റഡ് കോഡിനെ അത് ഉത്ഭവിച്ച DSL പ്രസ്താവനയുമായി മാപ്പ് ചെയ്യുന്നത് വളരെ ലളിതമാണ്. DSL-ന്റെ ഈ വശങ്ങൾക്ക് നന്ദി, കർണലുകൾ വികസിപ്പിക്കുമ്പോൾ ഞങ്ങൾ കൂടുതൽ ഉൽപ്പാദനക്ഷമതയുള്ളവരാണെന്ന് ഞങ്ങൾ കണ്ടെത്തി.

മോഡലുകൾ ഉടനീളം കർണലുകൾക്ക് നിരവധി സമാനതകൾ പങ്കിടാൻ കഴിയുമെങ്കിലും, പ്രകടനത്തെ സ്വാധീനിക്കാൻ കഴിയുന്ന വിവിധ കോൺഫിഗറേഷൻ നോബുകൾ അവയിൽ സജ്ജീകരിച്ചിരിക്കണം. പ്രസക്തമായ പാരാമീറ്ററുകൾ ഉൾപ്പെടുന്നു:

ഹിഡൻ ഡയമെൻഷൻ/ഹെഡ് ഡയമെൻഷൻ/അറ്റൻഷൻ ഹെഡുകളുടെ എണ്ണം

ഇൻപുട്ട്/ഔട്ട്പുട്ട് ഡാറ്റ തരങ്ങൾ

ക്വാണ്ടിൈസേഷൻ സ്കീം

കംപൈൽ-ടൈം കോൺസ്റ്റന്റുകൾ (ബയാസുകൾ, ആക്ടിവേഷൻ ഫംഗ്ഷനുകൾ മുതലായവ)

എഞ്ചിൻ, മോഡൽ, കേർണൽ ലൈബ്രറി, സേവനങ്ങൾ എന്നിവ ഉൾപ്പെടെയുള്ള ROSE ആർക്കിടെക്ചർ.

പ്രകടനത്തിനും ഡീബഗ്ഗിംഗിനും കംപൈൽ-ടൈം സ്പെഷ്യലൈസേഷൻ ഉപയോഗപ്രദമാണ്. പ്രത്യേകിച്ച് ഒരു പൈപ്പ്‌ലൈൻ പ്രൊസസ്സറിൽ, GPU-കളുടെ കാര്യത്തിലെന്നപോലെ, എപ്പോഴും എടുക്കുന്നതോ ഒരിക്കലും എടുക്കാത്തതോ ആയ ബ്രാഞ്ചിനേക്കാൾ നല്ലത് നിലവിലില്ലാത്ത ഒരു ബ്രാഞ്ചാണ്. കൂടാതെ, ഒരു നിർദ്ദിഷ്ട കോൺഫിഗറേഷൻ നോബുകളുടെ സെറ്റിനായി ഒരു കർണൽ കംപൈൽ ചെയ്യുന്നത് കുറഞ്ഞ PTX/SASS നിർദ്ദേശങ്ങൾക്ക് കാരണമാകുന്നു, അന്വേഷണ വേളയിൽ cuda-gdb ഉപയോഗിക്കുമ്പോൾ അവ വിശകലനം ചെയ്യാനും ഡീബഗ് ചെയ്യാനും എളുപ്പമാണ്.

നൽകിയിരിക്കുന്ന ബാച്ച് സൈസിനും കോൺഫിഗറേഷനും വേണ്ടി, DRAM-ലേക്കുള്ള പ്രതീക്ഷിക്കുന്ന ട്രാഫിക്കും അരിത്മെറ്റിക് പ്രവർത്തനങ്ങളുടെ എണ്ണവും കണക്കാക്കാം. തുടർന്ന്, ഒപ്റ്റിമൽ ഗ്രിഡ് അലോക്കേഷൻ, സിൻക്രൊണൈസേഷൻ, ചില ഹാർഡ്‌വെയർ സവിശേഷതകളുടെ ഉപയോഗം എന്നിവയ്ക്കായി കേസ്-ബൈ-കേസ് അടിസ്ഥാനത്തിൽ ഒപ്റ്റിമൽ തീരുമാനങ്ങൾ എടുക്കാം. സ്വീകാര്യതയില്ലാത്ത കംപൈൽ-ടൈം ഓവർഹെഡ് വരുത്താതെ തന്നെ ഡീപ് സ്പെഷ്യലൈസേഷനായി ഘടകങ്ങളെ നിർവചിക്കാനും വീണ്ടും ഉപയോഗിക്കാനും CuTeDSL നമ്മെ അനുവദിക്കുന്നു.

പൊതുവേ, ഇൻഫറൻസ് കർണലുകളുടെ പ്രകടനം രണ്ട് വ്യത്യസ്ത യൂസ് കേസുകൾക്കായി ഒപ്റ്റിമൈസ് ചെയ്യണം:

പ്രീഫിൽ: ചെറിയ ബാച്ച് സൈസ് (ചില ചങ്കുകൾ, സീക്വൻസുകൾ), ഓരോ സീക്വൻസിനും ധാരാളം ടോക്കണുകൾ. എംബെഡിംഗുകൾക്കും LLM-കളുടെ ഇൻപുട്ട് ടോക്കണുകൾക്കും ത്രൂപുട്ട് പരമാവധിയാക്കണം. പത്തുകൾ മുതൽ നൂറുകൾ വരെ മില്ലിസെക്കൻഡുകളിൽ ലേറ്റൻസി.

Decode: ചെറിയ-തിൽ-വലിയ batch sizes, ഓരോ sequence-ക്കും ഒരു ടോക്കൺ (speculative decoding ഇല്ല) മുതൽ കുറച്ച് ടോക്കണുകൾ വരെ (draft-target speculation, multi-token prediction). kernel launch latency-നോട് അടുത്ത latency.

GEMM-നും അറ്റൻഷനും വേണ്ടി, പ്രീഫിക്സ് ഘട്ടം വലിയൊരു എണ്ണം ടോക്കണുകളിൽ പ്രവർത്തിക്കുന്നു കൂടാതെ പ്രകടനം ഡിവൈസിന്റെ കമ്പ്യൂട്ട് ത്രൂപുട്ടാൽ പരിമിതപ്പെടുത്തപ്പെടുന്നു. ഡീകോഡിംഗിൽ ചെറിയ ബാച്ച് സൈസുകളിൽ ഒരേ പ്രവർത്തനം ഉൾപ്പെടുന്നു, ഡിവൈസിന്റെ മെമ്മറി ബാൻഡ്‌വിഡ്ത്ത് വഴി പരിമിതപ്പെടുത്തിക്കൊണ്ട്, ആക്ടിവേഷനുകളും വെയിറ്റുകളും DRAM-ൽ നിന്ന് സ്ട്രീമിംഗ് മൾട്ടിപ്രൊസസ്സറുകളിലേക്കും ടെൻസർ കോറുകളിലേക്കും ലോഡ് ചെയ്യുന്നതിന്റെ ചെലവാണ് എക്സിക്യൂഷൻ സമയങ്ങളിൽ ആധിപത്യം പുലർത്തുന്നത്. ബാച്ച് സൈസോ അല്ലെങ്കിൽ ഒരു ബാച്ചിലെ മൊത്തം ടോക്കണുകളുടെ എണ്ണമോ താരതമ്യേന വ്യക്തമായ സൂചകമാണെങ്കിൽ പോലും, പ്രീഫില്ലും ഡീകോഡും തമ്മിലുള്ള സ്വിച്ച്ഓവർ പരിധി കൂടുതൽ സൂക്ഷ്മമാണ്, കാരണം മെമ്മറി ബാൻഡ്‌വിഡ്ത്ത് മോഡലിന്റെ ഹിഡൻ ഡയമെൻഷനും ആക്ടിവേഷന്റെ ഡാറ്റ തരവും കൂടിയാണ് നിർണ്ണയിക്കുന്നത്. അൽഗോരിതം വ്യത്യാസങ്ങളോ ഗ്രിഡ്/ബ്ലോക്ക് അസൈൻമെന്റുകളോ അവയിലേതെങ്കിലും ഒന്നിന് മുൻതൂക്കം നൽകുകയാണെങ്കിൽ തൽഫലമായി കർണലുകൾ പ്രീഫിൽ, ഡീകോഡ് ഇംപ്ലിമെന്റേഷനുകളായി തിരിക്കപ്പെടുന്നു.

പ്രീഫിൽ, ഡീകോഡ് കർണലുകൾക്കിടയിൽ സാർവത്രികമായി വ്യത്യാസമുള്ള പാരാമീറ്റർ ഒരു ബ്ലോക്കിലെ വാർപ്പുകളുടെ എണ്ണമാണ് (ത്രെഡുകളുടെ എണ്ണം, ഓരോ വാർപ്പിനും 32 ത്രെഡുകൾ എന്ന വാർപ്പ് സൈസ് കൊണ്ട് ഹരിച്ചത്). രജിസ്റ്റർ അലോക്കേഷനെ സഹായിക്കാൻ കഴിയുമെന്നതിനാൽ, ഈ പാരാമീറ്റർ ഒരു കംപൈൽ-ടൈം കോസ്റ്റന്റ് ആക്കുന്നത് അഭികാമ്യമാണ്. CUDA-യിൽ, __launch_bounds__ വഴി ത്രെഡുകളുടെ എണ്ണം വ്യക്തമാക്കിക്കൊണ്ടും വാർപ്പുകളുടെ എണ്ണം ഒരു ടെംപ്ലേറ്റ് പാരാമീറ്ററായി നൽകിക്കൊണ്ടുമാണ് ഇത് സാധാരണയായി ചെയ്യുന്നത്:

cpp

template <size_t NUM_WARPS>
__global__ void 
__launch_bounds__(NUM_WARPS * WARP_SIZE, 1)
some_kernel(...) {

}

ഒരു ബ്ലോക്കിലെ എല്ലാ ത്രെഡുകൾക്കും ലഭ്യമായ രജിസ്റ്ററുകളുടെ എണ്ണം ഒരു ഹാർഡ്‌വെയർ പരിമിതിയാണ്, Hopper-ലും Blackwell-ലും ഇത് 64K ആണ്. ത്രെഡുകളുടെ എണ്ണം മുൻകൂട്ടി അറിയാമെങ്കിൽ, ഓരോ ത്രെഡിനും എത്ര രജിസ്റ്ററുകൾ അസൈൻ ചെയ്യാൻ കഴിയുമെന്ന് കംപൈലർക്ക് കൃത്യമായി അറിയാം, കൂടാതെ കൂടുതൽ ആക്രമണാത്മകമായി ഒപ്റ്റിമൈസ് ചെയ്ത കർണലുകൾ കംപൈൽ ചെയ്യുന്നതിൽ വിജയിക്കാനും കഴിയും. എന്നിരുന്നാലും, ഇതിന് കംപൈലേഷൻ-ടൈം ചെലവ് വരുന്നു, കാരണം ഇപ്പോൾ കർണൽ വിവിധ വാർപ്പ് വലുപ്പങ്ങൾക്കായി വീണ്ടും കംപൈൽ ചെയ്യേണ്ടതുണ്ട്. CuTeDSL ഉപയോഗിച്ച്, ഇത് ഇനി ഒരു പ്രശ്നമല്ല, കാരണം മുഴുവൻ കോൺഫിഗറേഷൻ സ്പേസിനുമുള്ള കംപൈലേഷൻ ചെലവ് ഞങ്ങൾ മുൻകൂട്ടി നൽകേണ്ടതില്ല. CUDA-യിലെ __latch_bounds__ സൂചിപ്പിക്കുന്ന ഒപ്റ്റിമൈസേഷൻ പരോക്ഷമായി നിർവഹിച്ചുകൊണ്ട്, ബ്ലോക്ക് സൈസ് ഒരു കോൺസ്റ്റന്റ് എക്സ്പ്രഷൻ ആയിരിക്കാനും DSL പ്രതീക്ഷിക്കുന്നു.

വാർപ്പുകളുടെ എണ്ണം ഒരു കോസ്റ്റന്റ് ആക്കുന്നതിന് പുറമേ, പ്രീഫില്ലും ഡീകോഡും തമ്മിൽ ഇത് വ്യത്യാസപ്പെടുത്തുന്നത് അഭികാമ്യമാണ്. ഒരു GPU-വിന്റെ മെമ്മറി ബാൻഡ്‌വിഡ്ത്ത് പൂർണ്ണമായി പ്രയോജനപ്പെടുത്തുന്നതിന്, ഒരു ബ്ലോക്കിന് ഉയർന്ന എണ്ണം വാർപ്പുകൾ വീതം എപ്പോഴും അസൈൻ ചെയ്യുന്നത് പര്യാപ്തമല്ല, കാരണം ലഭ്യമായ എല്ലാ സ്ട്രീമിംഗ് മൾട്ടിപ്രൊസസ്സറുകളും (SM-കൾ) പരമാവധി ഉപയോഗിച്ചാൽ മാത്രമേ പൂർണ്ണ ബാൻഡ്‌വിഡ്ത്ത് കൈവരിക്കാൻ കഴിയൂ. തൽഫലമായി, നൽകിയിരിക്കുന്ന ടോക്കണുകളുടെ എണ്ണത്തിനും അറിയപ്പെടുന്ന ഡാറ്റ തരങ്ങൾക്കും ഹിഡൻ ഡയമെൻഷനുകൾക്കും വേണ്ടി, ജനറേറ്റഡ് കോഡ് 128-ബിറ്റ് ലോഡ്/സ്റ്റോർ നിർദ്ദേശങ്ങൾ (Blackwell-ൽ 256-ബിറ്റ്) ഉപയോഗിക്കുന്നുവെന്ന് ഉറപ്പാക്കാൻ ഞങ്ങൾ ആദ്യം ഒരു അനുയോജ്യമായ പെർ-ത്രെഡ് വെക്ടർ ഡയമെൻഷൻ തിരഞ്ഞെടുക്കുന്നു (bfloat16-ന് 8, float32-ന് 4). വെക്ടർ വലുപ്പം അറിഞ്ഞുകൊണ്ട്, കർണലിന് ആവശ്യമായ വാർപ്പുകളുടെ എണ്ണം നമുക്ക് കണക്കാക്കാം. ഞങ്ങൾ ആദ്യം SM-കളുടെ എണ്ണം പരമാവധിയാക്കാൻ ശ്രമിക്കുന്നു: SM-കളേക്കാൾ കുറവ് വാർപ്പുകളാണ് ഉള്ളതെങ്കിൽ, ഒരു ബ്ലോക്കിന് ഒരൊറ്റ വാർപ്പ് തിരഞ്ഞെടുക്കപ്പെടുന്നു. ഇത് സാധാരണയായി കർണലിന്റെ ഡീകോഡ് വേരിയന്റാണ്. ആവശ്യത്തിന് വാർപ്പുകൾ ഉണ്ടെങ്കിൽ, വാർപ്പുകളുടെ എണ്ണം ക്രമേണ 8, 16 അല്ലെങ്കിൽ 32 ആയി വർദ്ധിപ്പിക്കുന്നു, അതിനുശേഷം പ്രകടനത്തിൽ യാതൊരു പുരോഗതിയും സാധാരണയായി കാണാത്ത ഒരു അപ്പർ ബൗണ്ട് ബെഞ്ച്മാർക്ക് ചെയ്യുന്നു.

Triton ഈ ഒപ്റ്റിമൈസേഷനുകളിൽ ചിലത് ഹുഡിന് കീഴിൽ നടത്തിയെങ്കിലും, ഉയർന്ന തലത്തിലുള്ള കംപൈലറുടെ പോരായ്മകളാൽ പരിമിതപ്പെടുത്തപ്പെടാതെ, CuTe ലേഔട്ട് ആൾജിബ്ര ശ്രദ്ധാപൂർവ്വം ഉപയോഗിച്ച് ഞങ്ങൾക്ക് സമാനമായ ഫലങ്ങൾ നേടാൻ കഴിയുമെന്ന് ഞങ്ങൾ കണ്ടെത്തി. പ്രത്യേകിച്ച് സബ്-വാർപ്പ് ഹോറിസോണ്ടൽ റിഡക്ഷനുകൾ ഉൾപ്പെട്ടിരിക്കുമ്പോൾ, ഒപ്റ്റിമൽ കോഡ് ജനറേറ്റുചെയ്യാൻ ആവശ്യമായ ഫൈൻ-ഗ്രേയ്ൻഡ് നിയന്ത്രണം CuTeDSL ഞങ്ങൾക്ക് വാഗ്ദാനം ചെയ്യുന്നുവെന്ന് ഞങ്ങൾ കണ്ടെത്തി.

ചില സാഹചര്യങ്ങളിൽ, കഴിയുന്നത്ര ബ്ലോക്കുകളിൽ ജോലി വിതരണം ചെയ്യുന്നതിന്, പങ്കെടുക്കുന്ന എല്ലാ ബ്ലോക്കുകളെയും സിൻക്രൊണൈസ് ചെയ്യാൻ cooperative_groups::this_grid().sync() ഉപയോഗിക്കാം. നിയന്ത്രണങ്ങൾ ബാധകമാണ്: എക്സിക്യൂഷൻ ആരംഭിച്ചുകഴിഞ്ഞാൽ ഒരു ബ്ലോക്കിന് ഒരു SM-ൽ നിന്ന് ഡി-ഷെഡ്യൂൾ ചെയ്യാൻ കഴിയില്ല എന്നതിനാൽ, മുന്നോട്ടുള്ള പുരോഗതി ഉറപ്പുനൽകുന്നതിന്, ഏതെങ്കിലും സമയത്ത് എല്ലാ ബ്ലോക്കുകൾക്കും സിൻക്രൊണൈസേഷനിൽ ഉൾപ്പെട്ടിരിക്കുന്ന ഗ്ലോബൽ മെമ്മറി ലൊക്കേഷനുകൾ അപ്ഡേറ്റ് ചെയ്യാനും പോൾ ചെയ്യാനും കഴിയുമെന്ന് ഉറപ്പാക്കാൻ, ബ്ലോക്കുകളുടെ എണ്ണം ഡിവൈസിൽ ലഭ്യമായ SM-കളുടെ കവിഞ്ഞുപോകാൻ പാടില്ല. ഈ നിയന്ത്രണം cudaLaunchCooperativeKernel കർശനമായി നടപ്പിലാക്കുന്നു. Hopper-ൽ ഒരു ഗ്രിഡ് ബാരിയറിന്റെ ചെലവ് ഏകദേശം 3us ആണ്, കൂടാതെ ചില സാഹചര്യങ്ങളിൽ ഒരു ഇംപ്ലിസിറ്റ് മെമ്മറി ബാരിയർ വഴി മുമ്പത്തെ എല്ലാ മെമ്മറി പ്രവർത്തനങ്ങളും ഫ്ലഷ് ചെയ്യുന്നതിന്റെ ചെലവും ഇതിൽ ഉൾപ്പെടുന്നു. മൊത്തത്തിൽ, ഇത് ഒരു കർണലിനെ രണ്ട് ലോഞ്ചുകളായി തിരിക്കുന്നതിന്റെ ഓവർഹെഡിന് തുല്യമാണ്, ഇത് പ്രകടനത്തെ സെൻസിറ്റീവായ ഒരു ഡിസൈൻ തീരുമാനത്തിലേക്ക് നയിക്കുന്നു. പ്രോഗ്രാമാറ്റിക് ഡിപെൻഡന്റ് ലോഞ്ചിന് (PDL) Hopper മുതൽ ലോഞ്ച് ഓവർഹെഡുകൾ കുറയ്ക്കാൻ കഴിയുമെന്ന വസ്തുത ഉണ്ടായിരുന്നിട്ടും, ഡീകോഡ് കർണലുകളെ സംബന്ധിച്ചിടത്തോളം, ഓവർഹെഡ് ഒഴിവാക്കുന്നത് മൂല്യവത്താണ്.

ജോലിഭാരങ്ങളിലൊന്നിന് SM-കളുടെ എണ്ണത്തേക്കാൾ കൂടുതൽ ബ്ലോക്കുകൾ ആവശ്യമാണെങ്കിൽ, കർണൽ ലോഞ്ചുകളുടെ ഇംപ്ലിസിറ്റ് ബാരിയറിനെ ആശ്രയിച്ചുകൊണ്ട് കർണൽ വിഭജിക്കപ്പെടണം. പ്രീഫില്ലിന്റെ കാര്യത്തിലാണ് ഇത് പ്രധാനമായും ഉള്ളത്. എല്ലാ ജോലിഭാരങ്ങൾക്കും അനുയോജ്യമായ ബ്ലോക്കുകളുടെ എണ്ണം ഗ്രിഡിൽ ഒതുങ്ങുകയാണെങ്കിൽ, ചെറിയൊരു എണ്ണം ടോക്കണുകളിൽ പ്രവർത്തിക്കുന്ന ഡീകോഡ് കർണലുകളുടെ കാര്യത്തിലെന്നപോലെ, ഒരു ഗ്രിഡ് ബാരിയർ ആണ് മുൻഗണന നൽകുന്നത്.

CuTeDSL-ൽ, ഈ തീരുമാനങ്ങൾ host-side launch code-ൽ എടുക്കാൻ കഴിയും, കൂടാതെ താഴെ കാണിച്ചിരിക്കുന്ന ഉദാഹരണത്തിൽ വിവരിച്ചതുപോലെ kernels-നെ ആവശ്യാനുസരണം compose ചെയ്യാനും grid barriers ഉപയോഗിച്ച് interleave ചെയ്യാനും കഴിയും. അത്തരത്തിലുള്ള hybrid strategy ഉപയോഗിച്ച്, decode kernels കുറഞ്ഞ latencies-നായി optimize ചെയ്യാനും prefill kernels പരമാവധി throughput-നായി optimize ചെയ്യാനും കഴിയും. CUDA-based solutions-നേക്കാൾ വ്യത്യസ്തമായി, ഈ logic കൂടുതൽ സംക്ഷിപ്തമായും compilation-time overhead കുറവുമായും നിർവചിക്കാം.

python

@cute.kernel
def dispatch_kernel_decode(...):
	route_tokens(...)
	grid_barrier(grid_barrier_ptr)
	copy_tokens_to_experts(...)

@cute.kernel
def dispatch_kernel_prefill_route(...):
	route_tokens(...)
	
@cute.kernel
def dispatch_kernel_prefill_copy(...):
	copy_tokens_to_experts(...)
	
@cute.jit
def dispatch(...):
	num_route_blocks = ...
	num_copy_blocks = ...
	
	if num_copy_blocks < NUM_SMS:
		dispatch_kernel_decode(...).launch(
			grid=(min(max(num_route_blocks, num_copy_blocks), NUM_SMS), 1, 1)
		)
	else:
		dispach_kernel_prefill_route(...).launch(
			grid=(num_route_blocks, 1, 1)
		)
		dispatch_kernel_prefill_copy(...).launch(
			grid=(num_copy_blocks, 1, 1)
		)

ഞങ്ങളുടെ MoE ഓൾ-ടു-ഓൾ കർണലുകളുടെ കാര്യത്തിൽ, താഴെ പറയുന്നതുപോലെ, ഗ്രിഡ് ബാരിയറുകളോടുകൂടിയോ അല്ലാതെയോ കർണലുകൾ നിർമ്മിക്കാനും കംപൈലേഷൻ സമയത്ത് അവ ഒപ്റ്റിമൈസ് ചെയ്യാനുമുള്ള വഴക്കം CuTeDSL വഴി കാര്യമായ പ്രകടന നേട്ടങ്ങൾ കൈവരിക്കാൻ ഞങ്ങളെ സഹായിച്ചു.

CuTeDSL-ൽ ഞങ്ങൾ നിർമ്മിച്ച ആദ്യത്തെ കർണൽ ഒരു പ്രത്യേക RMS നോം കർണൽ ആയിരുന്നു, ഇത് ചെറിയ ഡയമെൻഷനുകൾക്കായി ഒപ്റ്റിമൈസ് ചെയ്തതാണ്. സമീപകാല Qwen മോഡലുകളുടെ സെൽഫ്-അറ്റൻഷൻ ലെയറിൽ ഉപയോഗിക്കുന്ന QK നോം, ഓരോ ഹെഡിന്റെ അടിസ്ഥാനത്തിലും 64, 128 അല്ലെങ്കിൽ 256 എലമെന്റുകളിലാണ് പ്രയോഗിക്കുന്നത്. ആയിരക്കണക്കിന് കവിയുന്ന ഹിഡൻ ഡയമെൻഷനുകളിൽ പ്രവർത്തിക്കുന്ന, ഷെയർഡ് മെമ്മറിയും രജിസ്റ്ററുകളും വഴി ഒരു ഹോറിസോണ്ടൽ റിഡക്ഷൻ ആവശ്യമായ സാധാരണ RMS നോമിൽ നിന്ന് വ്യത്യസ്തമായി, QK നോമിന് വാർപ്പുകൾക്കിടയിൽ മാത്രം റിഡക്ഷൻ ആവശ്യമാണ്, ഒന്നിലധികം ഹെഡുകളെ ഒരൊറ്റ വാർപ്പിലേക്ക് ഒതുക്കാനുള്ള സാധ്യതയുമുണ്ട്. എല്ലാ ഹെഡ് ഡയമെൻഷനുകൾക്കും ഡാറ്റ തരങ്ങൾക്കും ഏറ്റവും മികച്ച ടോക്കൺ-ടു-വാർപ്പ് അസൈൻമെന്റുകൾ കണ്ടെത്താൻ കംപൈൽ സമയത്ത് ഞങ്ങൾക്ക് ആക്രമണാത്മകമായി ഒപ്റ്റിമൈസ് ചെയ്യാൻ കഴിയുന്ന ഒരു യൂസ് കേസ് ആയി ഇത് മാറി.

RMS നോം തികച്ചും ലളിതമായ ഒരു പ്രവർത്തനമാണ്, എങ്കിലും ഒരു വെക്ടറിന്റെ സ്ക്വയേർഡ് നോം കണക്കാക്കാൻ ഇതിന് ഒരു ഹോറിസോണ്ടൽ സം ആവശ്യമാണ്:

cpp

r = w * (x * ((x * x).sum(dim=-1) / x.shape[-1] + EPS))

അറ്റൻഷൻ ഹെഡുകൾ സാധാരണയായി ചെറുതും പെർ-ത്രെഡ് പ്രവർത്തനങ്ങൾ വളരെ ഉയർന്ന അളവിൽ വെക്ടറൈസ് ചെയ്തതുമായതിനാൽ, ഒരു വാർപ്പിന് മാത്രം സ്ക്വയർ നോം കണക്കാക്കാൻ കഴിയും. മെമ്മറി ബാൻഡ്‌വിഡ്ത്ത് പരമാവധിയാക്കുന്നതിന്, ഓരോ ത്രെഡും 128 ബിറ്റ് ഇൻപുട്ട് പ്രോസസ്സ് ചെയ്യണം, ഇത് 8 bfloat16 മൂല്യങ്ങളിലേക്ക് വിവർത്തനം ചെയ്യുന്നു. തൽഫലമായി, ഒരു ഹെഡ് പ്രോസസ്സ് ചെയ്യാൻ ഒരു ഹാഫ്-വാർപ്പ് പര്യാപ്തമാണ്, കൂടാതെ shuffle_sync_bfly നടത്തുന്ന ഒരു ഹോറിസോണ്ടൽ റിഡക്ഷൻ ഉപയോഗിച്ച് സ്ക്വയേർഡ് നോം നിർണ്ണയിക്കാൻ കഴിയും. ഹെഡ് ഡയമെൻഷനും ഡാറ്റ തരവും തുടങ്ങി, കംപൈലേഷൻ സമയത്ത് നമുക്ക് ഓരോ വാർപ്പിനും ഉള്ള ഹെഡുകളുടെ എണ്ണം നിർണ്ണയിക്കാനും ആർഗ്യുമെന്റുകളും റിഡക്ഷൻ റൗണ്ടുകളുടെ എണ്ണവും വ്യക്തമാക്കാനും കഴിയും.

സ്വയം-വെക്ടറൈസ്ഡ് ലോഡുകൾ/സ്റ്റോറുകളെ ആശ്രയിച്ചും വാർപ്പ് റിഡക്ഷനായി താഴ്ന്ന തലത്തിലുള്ള ആർക്കിടെക്ചർ-നിർദ്ദിഷ്ട പ്രവർത്തനങ്ങൾ ഉപയോഗിച്ചും താരതമ്യേന കുറഞ്ഞ വരി കോഡുകൾ ഉപയോഗിച്ച് ഞങ്ങൾക്ക് ഒരു പ്രത്യേക കർണൽ നിർമ്മിക്കാൻ കഴിയും. കർണൽ ഏത് ഹെഡിലാണ് (Q അല്ലെങ്കിൽ K) പ്രവർത്തിക്കേണ്ടത് എന്ന് നിർണ്ണയിക്കുകയും പ്രയോഗിക്കേണ്ട അനുബന്ധ വെയിറ്റുകൾ തിരിച്ചറിയുകയും ചെയ്യുന്നു. ഇത് ഹെഡിൽ നിന്ന് ഒരു മുഴുവൻ വെക്ടറും ലോഡ് ചെയ്യുകയും ആദ്യം പ്രാദേശികമായി സം ചെയ്യുകയും ചെയ്യുന്നു. തുടർന്ന്, ഓരോ ഹെഡിനുമുള്ള ത്രെഡുകളുടെ എണ്ണത്തെ അടിസ്ഥാനമാക്കി, സ്ക്വയേർഡ് നോം കണക്കാക്കപ്പെടുന്നു. അവസാനമായി, സ്വയം-വെക്ടറൈസ്ഡ് ആയ RMS നോം നിർണ്ണയിക്കപ്പെടുകയും മെമ്മറിയിലേക്ക് തിരികെ എഴുതപ്പെടുകയും ചെയ്യുന്നു. WEIGHT_BIAS, EPS പോലുള്ള മോഡൽ-നിർദ്ദിഷ്ട പാരാമീറ്ററുകൾ കോൺസ്റ്റന്റ്-ഫോൾഡ് ചെയ്യപ്പെടുന്നു.

python

@cute.kernel
def _qk_norm_kernel(
    q: cute.Tensor,
    k: cute.Tensor,
    q_weight: cute.Tensor,
    k_weight: cute.Tensor,
    WEIGHT_BIAS: cutlass.Constexpr[float],
    EPS: cutlass.Constexpr[float],
) -> None:
    # Find the block and thread indices.
    bidx, _, _ = cute.arch.block_idx()
    bdim, _, _ = cute.arch.block_dim()
    tidx, _, _ = cute.arch.thread_idx()

    _, num_q_heads, (threads_per_head, vecs_per_thread, vec_size) = q.shape
    _, num_k_heads, _ = k.shape
    head_dim = threads_per_head * vecs_per_thread * vec_size

    # Determine the thread and the vector to work on.
    vec_id = (bidx * bdim + tidx) // threads_per_head
    thread_id = tidx % threads_per_head

    # Determine the token and head indices.
    token = vec_id // (num_q_heads + num_k_heads)
    head_offset = vec_id % (num_q_heads + num_k_heads)
    is_q = head_offset < num_q_heads

    # Determine which tensor to work with.
    thread_coord = (thread_id, None, None)
    head_coord = (token, head_offset if is_q else head_offset - num_q_heads, thread_coord)

    # Load the head into registers.
    x_ref = q[head_coord] if is_q else k[head_coord]
    w = (q_weight if is_q else k_weight)[thread_coord].load()

    # Compute the squared sum locally.
    x = x_ref.load()
    x_f32 = x.to(cute.Float32)
    x_square = x_f32 * x_f32
    x_square_sum = x_square.reduce(cute.ReductionOp.ADD, init_val=0.0)

    # Sum up across the warp.
    mask = activemask()
    start = cute.arch.WARP_SIZE // (cute.arch.WARP_SIZE // threads_per_head * 2)
    for offset in [start >> i for i in range(5) if (start >> i) > 0]:
        x_square_sum += cute.arch.shuffle_sync_bfly(
            x_square_sum,
            offset=offset,
            mask=mask,
        )

    # RMS norm.
    rms = cute.math.rsqrt(x_square_sum / head_dim + EPS)
    x_norm = ((w + WEIGHT_BIAS) * (x * rms)).to(x.element_type)
    x_ref.store(x_norm)

സാധാരണ RMS നോം കർണലുമായി താരതമ്യപ്പെടുത്തുമ്പോൾ, ഈ പ്രത്യേക വേരിയന്റ് പ്രീഫില്ലിനും ഡീകോഡിനും 2-3x മികച്ച ത്രൂപുട്ടും ലേറ്റൻസിയും നൽകുന്നു.

MoE ഡിസ്പാച്ചും കമ്പൈനും ഇൻപുട്ട് ടോക്കണുകളെ അവ റൂട്ട് ചെയ്തിരിക്കുന്ന എക്സ്പെർട്ടുകളിലേക്ക് ഷഫിൾ ചെയ്യുന്നു, MoE GEMM കർണലുകൾ പ്രസക്തമായ മാട്രിക്സ് ഗുണനങ്ങൾ നടത്തിയതിന് ശേഷം വെയ്റ്റഡ് അക്യുമുലേഷൻ നടത്തുന്നു. വിവിധ ഉൽപ്പന്നങ്ങൾ വിവിധ വലുപ്പത്തിലുള്ള MoE മോഡലുകളെ ആശ്രയിക്കുന്നതിനാൽ, പ്രകടനം പ്രത്യേക താല്പര്യമുള്ള ഒന്നാണ്. വിന്യാസ തരങ്ങളെ അടിസ്ഥാനമാക്കി, ഒപ്റ്റിമൈസ് ചെയ്യാൻ ഞങ്ങൾ വിവിധ കർണലുകൾ വേർതിരിച്ചറിയുന്നു:

TP=1: ഒരൊറ്റ GPU-വിനുള്ള ലോക്കൽ കർണൽ. ഡിസ്പാച്ച് ടോക്കണുകളെ പ്രസക്തമായ എക്സ്പെർട്ടുകളിലേക്ക് ഷഫിൾ ചെയ്യുന്നു കൂടാതെ കമ്പൈൻ അവയെ അക്യുമുലേറ്റ് ചെയ്യുന്നു. H100-ൽ Qwen3 30B വിന്യാസങ്ങൾക്കായി പ്രധാനമായും ഉപയോഗിക്കുന്നു.

DP=1, TP ≠ 1: ടെൻസർ-പാരലൽ, സിംഗിൾ-നോഡ് കർണൽ. ഇൻപുട്ട് പകർപ്പുകൾ ഉണ്ടാക്കുന്നു, ഡിസ്പാച്ച് നിലവിലെ റാങ്കിനായുള്ള ടോക്കണുകൾ തിരഞ്ഞെടുക്കുന്നു, കൂടാതെ കമ്പൈൻ എല്ലാ റാങ്കുകളിലും ഉടനീളം അവയെ അക്യുമുലേറ്റ് ചെയ്യുന്നു. NVLink അല്ലെങ്കിൽ NVLS വഴിയുള്ള അക്യുമുലേഷൻ. വലിയ പ്രീഫില്ലറുകൾക്കായി ഉപയോഗിക്കുന്നു.

DP ≠ 1, NVLink: ഓപ്ഷണൽ ടെൻസർ പാരലലിസത്തോടുകൂടിയ ഡാറ്റ-പാരലൽ കർണൽ. ഓരോ റാങ്കും നിർദ്ദിഷ്ട എക്സ്പെർട്ടുകളെ ഹോസ്റ്റ് ചെയ്യുന്ന മറ്റെല്ലാ റാങ്കുകളിലേക്കും ടോക്കണുകൾ അയയ്ക്കുന്നു. റൂട്ടിംഗ് വിവരങ്ങൾ സംഗ്രഹിക്കപ്പെടുന്നു, NVIDIA NVLink വഴി ടോക്കണുകൾ കൈമാറ്റം ചെയ്യപ്പെടുന്നു. ഡാറ്റ-പാരലൽ H200 അല്ലെങ്കിൽ GB200 വിന്യാസങ്ങൾക്കായി ഉപയോഗിക്കുന്നു.

DP ≠ 1, InfiniBand: NVLink-നും NVIDIA ക്വാണ്ടം ഇൻഫിനിബണ്ടിനും ഇടയിൽ പ്രവർത്തിക്കുന്ന പീർ-ടു-പീർ ഡാറ്റ-പാരലൽ, ടെൻസർ-പാരലൽ കർണൽ. മൾട്ടി-നോഡ് Hopper വിന്യാസങ്ങൾക്കായി ഉപയോഗിക്കുന്നു.

റൂട്ടിംഗ് ടേബിളിനെ അടിസ്ഥാനമാക്കി ഓരോ എക്സ്പെർട്ടിനും അസൈൻ ചെയ്തിട്ടുള്ള ടോക്കണുകളുടെ എണ്ണം ലോക്കൽ കർണൽ കണക്കാക്കുകയും ഓരോ എക്സ്പെർട്ടിനും കീഴിലുള്ള ടോക്കണുകളുടെ ലിസ്റ്റിനുള്ളിൽ ഓരോ ടോക്കണിനും ഒരു അദ്വിതീയ ഇൻഡക്സ് നൽകുകയും ചെയ്യുന്നു. റൂട്ടിംഗ് ഘട്ടത്തിന് ശേഷം, ഒരു MoE GEMM കർണലിലേക്ക് കൈമാറുന്നതിനായി ഓപ്ഷണലായി പാഡ് ചെയ്ത ബഫറുകളിലേക്ക് ടോക്കണുകൾ പകർത്തപ്പെടുന്നു. ടോക്കൺ ലൊക്കേഷനുകൾ ഉപയോഗിച്ച്, കമ്പൈൻ കർണൽ പിന്നീട് രജിസ്റ്ററുകളിൽ ഒരു അക്യുമുലേഷൻ നടത്തുകയും അവസാന ടോക്കണുകളെ ഒരു ഔട്ട്പുട്ട് ബഫറിലേക്ക് എഴുതുകയും ചെയ്യുന്നു.

ഡിസ്പാച്ച് ഘട്ടം പ്രീഫില്ലും ഡീകോഡും തമ്മിൽ വ്യത്യാസപ്പെട്ടിരിക്കുന്നു: കുറച്ച് ടോക്കണുകൾക്ക്, വായിച്ചെടുക്കാൻ ചെറിയ റൂട്ടിംഗ് വിവരങ്ങളേ ഉള്ളൂ, അതിനാൽ ഞങ്ങൾ ഇത് ഒരൊറ്റ ബ്ലോക്കിൽ പ്രോസസ്സ് ചെയ്യുകയും ഗ്രിഡിന്റെ ബാക്കി ഭാഗത്തേക്ക് സിഗ്നൽ നൽകുകയും ചെയ്യുന്നു, വ്യക്തിഗത ബ്ലോക്കുകൾ ടോക്കണുകളുടെ ചങ്കുകൾ പകർത്തുന്നു. ഷെയർഡ് മെമ്മറിയിലെ കൗണ്ടറുകൾ ആറ്റോമിക്കായി ഇൻക്രിമെന്റ് ചെയ്തുകൊണ്ടാണ് ടോക്കൺ അസൈൻമെന്റ് നിർണ്ണയിക്കുന്നത്. പ്രീഫില്ലിന് കൂടുതൽ റൂട്ടിംഗ് വിവരങ്ങൾ പ്രോസസ്സ് ചെയ്യേണ്ടതുണ്ടെന്നതിനാൽ, ഒരു പ്രത്യേക കർണലിൽ റൂട്ടിംഗ് ഒന്നിലധികം ബ്ലോക്കുകളിലായി വിഭജിക്കപ്പെടുന്നു. ഷെയർഡ് മെമ്മറിയിലെ കൗണ്ടറുകൾ ഉപയോഗിച്ച് ഇൻഡക്സുകൾ ആദ്യം പ്രാദേശികമായി അസൈൻ ചെയ്യപ്പെടുന്നു, തുടർന്ന് ഓരോ ടോക്കണിന്റെയും സ്ഥാനം നിർണ്ണയിക്കാൻ ഗ്ലോബൽ മെമ്മറിയിലെ എല്ലാ ബ്ലോക്കുകളിലും അവ സംഗ്രഹിക്കപ്പെടുന്നു. ഗ്രിഡ്-വൈസ് ബാരിയറുകൾ ഉപയോഗിക്കാൻ കഴിയുന്നതിനായി ഈ കർണലിന്റെ ഗ്രിഡ് SM-കളുടെ എണ്ണത്തിലേക്ക് ക്ലാംപ് ചെയ്യപ്പെടുന്നു. ലഭ്യമായ മെമ്മറി ബാൻഡ്‌വിഡ്ത്ത് പൂർണ്ണമായി പ്രയോജനപ്പെടുത്തുന്നതിനായി വളരെ വലിയൊരു എണ്ണം ബ്ലോക്കുകളിലും വാർപ്പുകളിലും ലോഞ്ച് ചെയ്തുകൊണ്ട് തുടർന്നുള്ള ഒരു കർണൽ പകർപ്പുകൾ നിർവഹിക്കുന്നു.

കമ്പൈൻ കർണൽ താരതമ്യേന ലളിതമാണ്, കാരണം ഇത് വെക്ടറൈസ്ഡ് അക്യുമുലേഷൻ നടത്തുന്നു, ഇത് റൂട്ടിംഗ് വിവരങ്ങളെയും ഡിസ്പാച്ച് കണക്കാക്കിയ ഓഫ്സെറ്റുകളെയും ആശ്രയിക്കുന്നു. എല്ലാ ഡാറ്റ തരങ്ങൾക്കും ഹിഡൻ ഡയമെൻഷനുകൾക്കും കാര്യക്ഷമമായ കോഡിലേക്ക് കംപൈൽ ചെയ്യുന്ന ലളിതമായ CuTe ടെൻസറുകൾ വഴി കമ്പ്യൂട്ടേഷൻ പ്രകടിപ്പിക്കാൻ കഴിയുമെന്നതിനാൽ CuTeDSL ഇവിടെ വിലമതിക്കാനാവാത്തതാണെന്ന് തെളിയിക്കപ്പെട്ടു.

ഒരു ടെൻസർ-പാരലൽ വിന്യാസത്തിൽ, റൂട്ടിംഗ് വിവരങ്ങളും ഇൻപുട്ട് ടോക്കണുകളും പകർപ്പുകളുണ്ടാക്കപ്പെടുന്നു. നിലവിലെ റാങ്കിൽ സ്ഥിതി ചെയ്യുന്ന എക്സ്പെർട്ടുകൾക്ക് അസൈൻ ചെയ്തിട്ടുള്ള ടോക്കണുകൾ മാത്രം റൂട്ട് ചെയ്യാനും തിരഞ്ഞെടുക്കാനുമുള്ള ചെറിയൊരു മാറ്റത്തോടെ ലോക്കൽ ഡിസ്പാച്ച് കർണൽ പൂർണ്ണമായി വീണ്ടും ഉപയോഗിക്കാൻ കഴിയും. ചില പാരാമീറ്ററുകളും കംപൈൽ-ടൈം സ്പെഷ്യലൈസേഷനും ഉപയോഗിച്ച്, ലോക്കൽ ഡിസ്പാച്ച് കർണലുകൾ വീണ്ടും ഉപയോഗിക്കപ്പെടുന്നു.

റാങ്കുകളുടെ എണ്ണം സാധാരണയായി ഒരു ടോക്കൺ റൂട്ട് ചെയ്യപ്പെടുന്ന എക്സ്പെർട്ടുകളുടെ എണ്ണത്തേക്കാൾ കുറവായതിനാൽ, ഓരോ ടോക്കണിനും അതിന്റെ പകർപ്പുകളിൽ കുറഞ്ഞത് ഒരെണ്ണമെങ്കിലും ഓരോ റാങ്കിലേക്കും റൂട്ട് ചെയ്യപ്പെട്ടിട്ടുണ്ടെന്ന് ശരാശരി കേസിൽ അനുമാനിക്കാം. ഈ സാഹചര്യത്തിൽ, പങ്കെടുക്കുന്ന റാങ്കുകളിലുടനീളം തുക കാണുന്നതിന് ഒരു ഓൾ-റെഡ്യൂസ് തുടർന്ന് വന്നുകൊണ്ട്, ഒരു ടോക്കണിന്റെ ലോക്കൽ പകർപ്പുകൾ അക്യുമുലേറ്റ് ചെയ്യാൻ ലോക്കൽ കമ്പൈൻ കർണൽ വീണ്ടും ഉപയോഗിക്കാം. പ്രീഫിൽ കർണലിന്റെ കാര്യത്തിൽ, ഞങ്ങൾ പ്രത്യേക കർണലുകൾ ലോഞ്ച് ചെയ്യുന്നു, എന്നിരുന്നാലും ഡീകോഡ് രണ്ട് പ്രവർത്തനങ്ങളെയും ഒരൊറ്റ ലോഞ്ചിലേക്ക് സംയോജിപ്പിക്കുന്നു.

ഓൾ-റെഡ്യൂസിന് വേണ്ടി, രണ്ടോ അതിലധികമോ റാങ്കുകൾ മാത്രമാണ് പങ്കെടുക്കുന്നതെങ്കിൽ ഞങ്ങൾ പ്ലെയിൻ NVLink ഉപയോഗിക്കുന്നു, ത്രൂപുട്ടിൽ 10-20% വർദ്ധനവിനായി 4 അല്ലെങ്കിൽ അതിൽ കൂടുതൽ റാങ്കുകൾ ഉൾപ്പെടുമ്പോൾ NVLS-ലേക്ക് മാറുമ്പോൾ. DSL-ന് നന്ദി, MoE-യ്ക്കും സാധാരണ ഓൾ-റെഡ്യൂസ് കർണലുകൾക്കുമിടയിൽ ഭൂരിഭാഗം റിഡക്ഷൻ ലോജിക്കും പീർ ബഫർ കൈമാറ്റങ്ങളും ഞങ്ങൾക്ക് പങ്കിടാൻ കഴിയും. ഹാർഡ്‌വെയറിന്റെ മേൽ അടുത്ത നിയന്ത്രണം പ്രയോജനപ്രദമാകുന്ന ഒരു സാഹചര്യമാണിത്, കാരണം NVLS-ന് multimem.ld_reduce, multimem.st എന്നിവയ്ക്ക് മുകളിലുള്ള DSL പ്രവർത്തനങ്ങൾ ആവശ്യമാണ്. റിഡക്ഷന്റെ ഇംപ്ലിമെന്റേഷൻ തന്നെ NCCL ടെസ്റ്റുകളിൽ നിന്നുള്ള all_reduce ഉദാഹരണത്തെയാണ് വലിയ തോതിൽ പിന്തുടരുന്നത്.

നിരവധി ഉപകരണങ്ങൾക്കായി ടോക്കണുകൾ സംഗ്രഹിക്കാൻ ഡാറ്റ പാരലൽ കർണലുകൾ NVLink-നെ ആശ്രയിക്കുന്നു. ഡിസ്പാച്ച് ഘട്ടത്തിൽ, ടോക്കണുകൾ കൈമാറുന്നതിന് മുമ്പ്, അയയ്ക്കുന്ന ഓരോ ടോക്കണിനും ഓരോ റാങ്കിലെയും ടാർഗെറ്റ് ഓഫ്സെറ്റ് തിരിച്ചറിയാൻ പീറുകൾ ആദ്യം പ്രാദേശികമായി റൂട്ട് ചെയ്യുന്നു. തുടർന്ന്, വൈരുദ്ധ്യങ്ങൾ ഒഴിവാക്കാൻ ഓരോ അയച്ചയാൾക്കുമായി അലോക്കേറ്റ് ചെയ്ത റിസീവ് ബഫറുകളിലേക്ക് അവർ ടോക്കണുകൾ പകർത്താൻ തുടങ്ങുന്നു. പൂർത്തിയായെന്ന് സൂചിപ്പിക്കാൻ പെർ-ടോക്കൺ റിലീസ്-സ്റ്റോർ റൈറ്റുകൾ ഉപയോഗിക്കുന്നു. സ്വീകരിക്കുന്ന ഭാഗത്ത്, MoE GEMM-നായി എല്ലാ ടോക്കണുകളെയും ഒരൊറ്റ ബഫറിലേക്ക് ഒതുക്കാൻ കഴിയുന്നതിനായി റൂട്ടിംഗ് വിവരങ്ങൾ സംഗ്രഹിക്കപ്പെടുന്നു. എഴുത്തുകൾ പൂർത്തിയാകുന്നതിനായി റിസീവറുകൾ അക്വയർ-ലോഡ് ഉപയോഗിച്ച് പീർ എഴുതിയ ഫ്ലാഗ് പോൾ ചെയ്യുന്നു.

കമ്പൈൻ വശത്ത്, ഒരു പെർ-ടോക്കൺ ഫ്ലാഗ് സ്വീകാര്യത സൂചിപ്പിച്ചുകൊണ്ട്, ഉറവിട റാങ്കുകളിലേക്ക് NVLink വഴി ടോക്കണുകൾ പകർത്തപ്പെടുന്നു. വിലാസം നൽകുന്നതിലെ ചെറിയ വ്യത്യാസങ്ങളും ചില അധിക സിൻക്രൊണൈസേഷൻ ലോജിക്കും ഒഴിച്ചാൽ കമ്പൈൻ കർണൽ ലോക്കൽ, സിംഗിൾ-നോഡ് വേരിയന്റുകൾക്ക് സമാനമാണ്.

ഈ കർണൽ നടപ്പിലാക്കുന്നതിനായി സിസ്റ്റം, ഡിവൈസ് തലങ്ങളിൽ സ്കോപ്പ് ചെയ്തിട്ടുള്ള റിലീസ്-അക്വയർ ലോഡുകളും സ്റ്റോറുകളും എക്സ്പോസ് ചെയ്യുന്നതിനായി ഞങ്ങൾ PTX റാപ്പറുകളുടെ ഒരു ലൈബ്രറി നിർമ്മിച്ചു. സിൻക്രൊണൈസേഷൻ ഫ്ലാഗുകളുടെ ഗ്രാനുലാരിറ്റി മെച്ചപ്പെടുത്തുന്നതിലൂടെയും NVLink വഴി ടോക്കണുകൾ അയയ്ക്കാനും സ്വീകരിക്കാനുമായി കൂടുതൽ ബ്ലോക്കുകൾ അലോക്കേറ്റ് ചെയ്യുന്നതിലൂടെയും, ഞങ്ങളുടെ മുൻകാല CUDA കർണലുകളെ അപേക്ഷിച്ച് ഞങ്ങൾ 2x മെച്ചപ്പെടുത്തലുകൾ കൈവരിച്ചു.

ഞങ്ങളുടെ പീർ-ടു-പീർ കർണലുകൾ ഇതിനുമുമ്പ് ConnectX-7 ഇൻഫിനിബഡ് അഡാപ്റ്ററുകളുമായി ജോടിയാക്കിയ Hopper GPU-കളിൽ അത്യാധുനിക പ്രകടനം കൈവരിച്ചിരുന്നു. പ്രകടനത്തിൽ കൂടുതൽ നേട്ടങ്ങൾക്കായി ഞങ്ങൾ അതിനുശേഷം കമ്പൈൻ കർണലുകൾ CuTeDSL-ലേക്ക് പോർട്ട് ചെയ്തിട്ടുണ്ട്. CUDA പതിപ്പുകൾ ഇതിനകം തന്നെ കംപൈൽ-ടൈം സ്പെഷ്യലൈസേഷനെ വലിയ തോതിൽ ആശ്രയിച്ചിരുന്നതിനാൽ, ഞങ്ങളുടെ ശേഖരത്തിലെ കംപൈലേഷൻ-ടൈം തടസ്സങ്ങളായി അവ തെളിഞ്ഞു, ചില പാരാമീറ്ററുകൾ ഒഴിവാക്കാൻ ഞങ്ങളെ നിർബന്ധിതരാക്കി. JIT-കംപൈൽ ചെയ്ത DSL-ലേക്ക് മാറുന്നതിലൂടെ, ഞങ്ങൾ കൂടുതൽ പാരാമീറ്ററുകൾക്കായി സ്പെഷ്യലൈസ് ചെയ്യുകയും ഗ്രിഡ് അലോക്കേഷനുകൾ ട്വീക്ക് ചെയ്യുകയും ചെയ്തു, ലേറ്റൻസിയിൽ ഏകദേശം 3% വർദ്ധനവിനായി പ്രീഫില്ലും ഡീകോഡും തമ്മിൽ കൂടുതൽ വ്യത്യാസങ്ങൾ അവതരിപ്പിച്ചു.

ROSE കർണലുകളുടെ കാഴ്ചപ്പാടിൽ, Blackwell പിന്തുണ രണ്ട് വിഭാഗങ്ങളായി പെടുന്നു: ആർക്കിടെക്ചർ-അഗ്നോസ്റ്റിക്, ആർക്കിടെക്ചർ-നിർദ്ദിഷ്ട കർണലുകൾ. ടെൻസർ കോറുകളെ ആശ്രയിക്കുന്ന ആർക്കിടെക്ചർ-നിർദ്ദിഷ്ട കർണലുകൾക്ക് പൊരുത്തപ്പെടാൻ വ്യക്തമായ ജോലി ആവശ്യമാണ്. എന്നിരുന്നാലും, DSL-ന്റെ ഉപയോഗത്തിന് നന്ദി, ആർക്കിടെക്ചർ-അഗ്നോസ്റ്റിക് കർണലുകൾ അനായാസമായി പൊരുത്തപ്പെടുന്നു.

ഒരു Blackwell GPU-യ്ക്ക് Hopper-ന്റെ 132-നോട് താരതമ്യത്തിൽ പരമാവധി 160 SMs ഉണ്ട്. ഭൂരിഭാഗം kernels-ക്കായി, ഇത് ഒരു compile-time constant ആണ്, കൂടാതെ host-side launcher code-ലുള്ള logic, പ്രവർത്തിപ്പിക്കേണ്ട ഏറ്റവും അനുയോജ്യമായ (prefill vs decode) kernel നിർണ്ണയിക്കുമ്പോൾ ഇതിനെ പരിഗണിക്കുന്നു. കൂടാതെ, Blackwell 256-bit load instructions-ഉം അവതരിപ്പിക്കുന്നു, ഇത് vector width (CuTe tensors-ന്റെ അവസാന dimension) വർധിപ്പിച്ച് memory bandwidth utilization മെച്ചപ്പെടുത്താൻ എളുപ്പത്തിൽ emit ചെയ്യാൻ കഴിയും.

ഗണ്യമായ ഒരു എണ്ണം ക്വാണ്ടിൈസേഷൻ സ്കീമുകൾക്ക് (MXFP8, MXFP4, NVFP4) Blackwell ഹാർഡ്‌വെയർ പിന്തുണ അവതരിപ്പിക്കുന്നു. FP8 ബ്ലോക്ക്-വൈസ് സ്കെയിലിംഗിനായി നിലവിലുള്ള കർണലുകൾ ട്വീക്ക് ചെയ്തുകൊണ്ടാണ് ഇവയ്ക്കുള്ള പിന്തുണ പ്രധാനമായും നടപ്പിലാക്കിയത്. കംപൈൽ-ടൈം കോൺഫിഗറേഷൻ സ്പേസ് പൊട്ടിത്തെറിക്കാതെ തന്നെ പുതിയ മോഡലുകൾ പിന്തുണയ്ക്കാൻ കഴിയുമെന്നതിനാൽ JIT പിന്തുണ ഇവിടെ നിർണ്ണായകമാണ്. മുൻകാലങ്ങളിൽ, ചുരുക്കം ചില എലമെന്റുകളിൽ ഉടനീളം, ഒരു വാർപ്പിന്റെ പകുതിയോ കാൽഭാഗമോ മീതെ ഹോറിസോണ്ടൽ റിഡക്ഷനുകൾ പ്രകടിപ്പിക്കുന്നതിൽ Triton അസ്വാഭാവികമാണെന്ന് ഞങ്ങൾ കണ്ടെത്തി. NVFP4, MXFP8/MXFP4 എന്നിവയ്ക്കായി യഥാക്രമം 16 അല്ലെങ്കിൽ 32 എലമെന്റുകളിൽ ഉടനീളം ഒരു കേവല പരമാവധി മൂല്യം കണക്കാക്കുമ്പോൾ ഒപ്റ്റിമൈസ് ചെയ്ത പ്രകടനം കൈവരിക്കുന്നതിന് റിഡക്ഷനുകളുടെ മേലുള്ള മികച്ച ഗ്രാനുലാർ നിയന്ത്രണം നിർണ്ണായകമായിരുന്നു.

താരതമ്യേന പുതിയ ഒരു ഭാഷ എന്ന നിലയിൽ, പ്രീ-ട്രെയിനിംഗ് വേളയിൽ CuTeDSL കോഡുമായി LLM-കൾക്ക് ചെറിയ എക്സ്പോഷർ മാത്രമേ ഉണ്ടായിരുന്നുള്ളൂ. കൂടാതെ, ആക്ടിവേഷനുകൾക്കും നോമുകൾക്കുമായി ഒപ്റ്റിമൈസ് ചെയ്ത കർണലുകളോടുകൂടിയ ഞങ്ങളുടെ നിലവിലെ കോഡിന്റെ ശൈലി, TMA-യ്ക്കും ആർക്കിടെക്ചർ-നിർദ്ദിഷ്ട മാട്രിക്സ് ഗുണന നിർദ്ദേശങ്ങൾക്കുമായി ഉയർന്ന അളവിൽ ഒപ്റ്റിമൈസ് ചെയ്ത നിലവിലുള്ള CUTLASS അല്ലെങ്കിൽ Quack ഉദാഹരണങ്ങളിൽ നിന്ന് അൽപ്പം വ്യത്യസ്തമായിരുന്നു. കർണൽ എഴുതുന്നതിൽ LLM-കൾക്ക് സഹായിക്കാൻ കഴിയുന്ന ഒരു ഘട്ടത്തിലെത്തുന്നത് ക്രമാനുഗതമായ ഒരു പ്രക്രിയയായിരുന്നു.

കർണൽ പരിവർത്തനത്തിനായി ഏജന്റുകളെ ആശ്രയിക്കുന്നതിന് മുമ്പ്, സ്ഥിരമായ പാറ്റേണുകൾ പിന്തുടർന്ന് ഞങ്ങൾ കർണലുകളുടെയും യൂട്ടിലിറ്റി ഫംഗ്ഷനുകളുടെയും ഒരു പ്രാരംഘ കോർപ്പസ് സ്വയം നിർമ്മിച്ചു. കോർപ്പസ് പുരോഗമിച്ചപ്പോൾ, CUDA-യിൽ നിന്ന് CuTeDSL-ലേക്ക് ഞങ്ങൾക്ക് ഏകദേശ വിവർത്തനങ്ങൾ നടത്താൻ കഴിഞ്ഞു, അതിനുശേഷം കാര്യമായ മാനുവൽ റീവർക്ക് ആവശ്യമായി വന്നു. പ്രത്യേകിച്ചും, പോയിന്റർ-ടു-ടെൻസർ പരിവർത്തനങ്ങളിലും വെക്ടറൈസേഷനിലും മോഡലുകൾ ബുദ്ധിമുട്ടുന്നതായി ഞങ്ങൾ കണ്ടെത്തി, ഒരു ടെസ്റ്റ് കംപൈൽ ചെയ്യാൻ ലക്ഷ്യമിടുമ്പോൾ ടൈപ്പ് പിഴവുകൾ പരിഹരിക്കാൻ ശ്രമിച്ചു കുടുങ്ങിപ്പോയി. കർണലുകളുടെ ഗ്രിഡ്-അധിഷ്ഠിത സ്വഭാവത്തെക്കുറിച്ച് അവയ്ക്ക് ചെറിയ ധാരണയേ ഉണ്ടായിരുന്നുള്ളൂ, ത്രെഡും ഗ്രിഡ് സിൻക്രൊണൈസേഷനും ആശയക്കുഴപ്പത്തിലാക്കുന്നു. ഏറ്റവും നിരാശാജനകമെന്നോണം, ഞങ്ങളുടെ സ്വന്തം CuTeDSL പകരക്കാർക്ക് എതിരായി ബിൽഡ് ചെയ്യുന്നതിന് പകരം CUDA-നിർദ്ദിഷ്ട ലൈബ്രറികളുടെ നിലവിലില്ലാത്ത വിവർത്തനങ്ങൾ ഭ്രമാത്മകം ചെയ്യാനുള്ള പ്രവണത ഉണ്ടായിരുന്നു. കാലക്രമേണ, CuTeDSL കർണലുകളുടെ കോർപ്പസ് വളർന്നപ്പോൾ, സ്കെലട്ടൺ കോഡിന്റെ ഭൂരിഭാഗവും നിർമ്മിക്കാനും ശൂന്യമായ ഭാഗങ്ങൾ പൂരിപ്പിക്കാനും LLM-കൾക്ക് കഴിയുമെന്ന് ഞങ്ങൾ കണ്ടെത്തി, ഇത് ഉയർന്ന തലത്തിലുള്ള രൂപകൽപ്പനയിൽ ശ്രദ്ധ കേന്ദ്രീകരിക്കാൻ ഞങ്ങളെ അനുവദിച്ചു.

CUDA-യുമായി താരതമ്യപ്പെടുത്തുമ്പോൾ, CuTeDSL വളരെ മികച്ച കംപൈൽ-ടൈം പ്രകടനം നൽകുന്നുണ്ടെന്ന് ഞങ്ങൾ കണ്ടെത്തി, കൂടുതൽ പാരാമീറ്ററുകൾക്കായി സ്പെഷ്യലൈസ് ചെയ്യാനും പ്രകടന മെച്ചപ്പെടുത്തലുകൾ കൈവരിക്കാനും ഇത് ഞങ്ങളെ അനുവദിക്കുന്നു. Triton-മായി താരതമ്യപ്പെടുത്തുമ്പോൾ ചില അമർബ്സ്ട്രാക്ഷനുകളുടെ അഭാവം കാരണം കർണലുകൾ ചെറുതായി കൂടുതൽ വാചാലമാണെങ്കിലും, ഹാർഡ്‌വെയറിന്റെ മേലുള്ള മികച്ച ഗ്രാനുലാർ നിയന്ത്രണം വേഗതയേറിയ കർണലുകൾ നിർമ്മിക്കാൻ ഞങ്ങളെ അനുവദിക്കുന്നതിനാൽ സങ്കീർണ്ണത ന്യായമാണെന്ന് ഞങ്ങൾ കണ്ടെത്തി.

എഴുതുന്ന സമയത്ത്, CUTLASS C++, CuTeDSL കർണലുകൾ എന്നിവ തമ്മിലുള്ള പ്രകടന അന്തരം കുറവാണ്, ഉയർന്ന പ്രകടനമുള്ള GEMM കർണലുകളെ ഡിവൈസ്-തല ആശയവിനിമയം, നോർമലൈസേഷൻ അല്ലെങ്കിൽ ആക്ടിവേഷനുകൾ എന്നിവയുമായി സംയോജിപ്പിക്കുന്നതിനും സംയോജിപ്പിക്കുന്നതിനുമുള്ള ആവേശകരമായ അവസരങ്ങൾ ഇത് തുറന്നുനൽകുന്നു. ഈ DSL-ൽ പ്രതിജ്ഞാബദ്ധരാകുന്നതിലൂടെ, എംബെഡിംഗുകൾ മുതൽ LLM-കൾ വരെയുള്ള ഞങ്ങളുടെ ജോലിഭാരങ്ങളിൽ ഉടനീളം ത്രൂപുട്ട് മെച്ചപ്പെടുത്തുന്നതിന് ഞങ്ങളുടെ ഉപകരണങ്ങളിൽ നിന്ന് കൂടുതൽ പ്രകടനം വേർതിരിച്ചെടുക്കാൻ ഞങ്ങൾക്ക് കഴിയും.

https://github.com/Dao-AILab/quack

https://github.com/NVIDIA/cutlass

https://github.com/NVIDIA/nccl-tests

https://docs.nvidia.com/cuda/parallel-thread-execution/

BibTeX

@online{perplexity2026CutedslAtPerplexity,
  author = {{Perplexity} and Nandor Licker and Bex Brooker and Jesse Dwyer and Phi Hoang},
  title = {{CuTeDSL at Perplexity}},
  date = {2026-05-06},
  url = {https://www.perplexity.ai/hub/blog/cutedsl-at-perplexity}
}

Perplexity AI-ലെ ഇൻഫറൻസ്

എന്തുകൊണ്ട് CuTeDSL?

ഡെവലപ്പർ അനുഭവം

വലിയ കോൺഫിഗറേഷൻ സ്പേസ്

പ്രീഫില്ലും ഡീകോഡ് സ്പെഷ്യലൈസേഷനും

ഗ്രിഡ് സിൻക്രൊണൈസേഷൻ

CuTeDSL കർണലുകൾ

QK നോം

മിക്സ്ചർ-ഓഫ്-എക്സ്പെർട്സ് ഡിസ്പാച്ച്/കമ്പൈൻ

ലോക്കൽ

സിംഗിൾ-നോഡ്

ഡാറ്റ പാരലൽ

പീർ-ടു-പീർ

Blackwell പിന്തുണ

AI-അസിസ്റ്റഡ് കർണൽ ജനറേഷൻ

ഉപസംഹാരവും ഭാവി ജോലികളും

റഫറൻസുകൾ

ഉദ്ധരണി