ലാർജ് ലാംഗ്വേജ് മോഡലുകൾ (LLM) ഒരു ഫോണിൽ പ്രവർത്തിപ്പിക്കുക എന്നത് ഇനി വെറുമൊരു ഗവേഷണ പരീക്ഷണമല്ല. അത് പ്രായോഗികമായ ഒരു യാഥാർത്ഥ്യമാണ്. എന്നാൽ, ഒരു ചെറിയ ഡെമോയിൽ നിന്ന് യഥാർത്ഥ ഉൽപ്പന്നത്തിലേക്ക് നിങ്ങൾ മാറുമ്പോൾ, ലേറ്റൻസി (latency) വീണ്ടും പ്രശ്നമായി മാറുന്നു. നിങ്ങൾ മോഡൽ ട്രിം ചെയ്യുകയും വെയ്റ്റുകൾ ക്വാണ്ടൈസ് (quantize) ചെയ്യുകയും ചെയ്തിട്ടും പ്രീഫിൽ (prefill) ഘട്ടം മന്ദഗതിയിലായി തുടരുന്നു. ടോക്കണുകൾ തടസ്സപ്പെടുന്നു. UI ഫ്രീസ് ആകുന്നു. ഇതിന്റെ കുറ്റം പലപ്പോഴും ശരിയായ സ്ഥലത്ത് എത്തുന്നില്ല.

ആൻഡ്രോയിഡ് ഉപകരണങ്ങളിൽ, LLM പ്രീഫിൽ സമയത്ത് കമ്പ്യൂട്ട് (compute) എന്നത് ഒരിക്കലും ഒരു തടസ്സമാകാറില്ല. പകരം മെമ്മറി ബാൻഡ്‌വിഡ്ത്ത് (memory bandwidth) ആണ് തടസ്സം. ആധുനിക ഫ്ലാഗ്ഷിപ്പ് SoCs-ൽ ശക്തമായ GPU, NPU കോറുകൾ ഉണ്ട്; ഇവയ്ക്ക് മെമ്മറി സബ്സിസ്റ്റം നൽകുന്നതിനേക്കാൾ വളരെ വേഗത്തിൽ ഗണിതക്രിയകൾ (arithmetic) ചെയ്യാൻ കഴിയും. ഒരു സാധാരണ അറ്റൻഷൻ ഇംപ്ലിമെന്റേഷൻ പ്രൊഫൈൽ ചെയ്യുമ്പോൾ, എക്സിക്യൂഷൻ യൂണിറ്റുകൾ പൂർണ്ണമായി ഉപയോഗിക്കപ്പെടുന്നില്ല എന്ന് കാണാം. അവ കാത്തിരിക്കുകയാണ്. DRAM-നായി കാത്തിരിക്കുന്നു.

നിങ്ങളുടെ ക്വാണ്ടൈസ് ചെയ്ത മോഡൽ ഇപ്പോഴും സാവധാനത്തിലായി തോന്നുന്നത് എന്തുകൊണ്ട്?

ഓൺ-ഡിവൈസ് ഇൻഫറൻസിനായി ക്വാണ്ടൈസേഷൻ (Quantization) ഇപ്പോൾ ആദ്യപടിയായി മാറിയിരിക്കുന്നു. FP16-ൽ നിന്ന് INT8-ലേക്ക് വെയ്റ്റുകൾ കുറയ്ക്കുന്നത് മോഡലിന്റെ വലിപ്പം പകുതിയാക്കുകയും സ്റ്റോറേജ് കുറയ്ക്കുകയും ചെയ്യുന്നു. ഇത് സഹായിക്കുമെങ്കിലും, അറ്റൻഷൻ ലെയറിലെ ലേറ്റൻസി പരിഹരിക്കാൻ ഇതിന് കഴിയില്ല. കാരണം ലളിതമാണ്: ക്വാണ്ടൈസേഷൻ നിങ്ങൾ സൂക്ഷിക്കുന്ന ഡാറ്റയുടെ അളവ് കുറയ്ക്കുന്നു, എന്നാൽ അറ്റൻഷൻ മെക്കാനിസം നടത്തുന്ന മെമ്മറി ട്രാൻസാക്ഷനുകളുടെ എണ്ണം അത് കുറയ്ക്കുന്നില്ല.

പാഠപുസ്തക രീതിയിൽ നടപ്പിലാക്കിയ ഒരു സാധാരണ മൾട്ടി-ഹെഡ് അറ്റൻഷൻ ലെയർ, ഓരോ ലെയറിനും DRAM-ലേക്ക് മൂന്ന് തവണ പൂർണ്ണമായി റൗണ്ട് ട്രിപ്പുകൾ നടത്തുന്നു. ക്വറി (Query), കീ (Key), വാല്യൂ (Value) മാട്രിക്സുകൾ മെയിൻ മെമ്മറിയിൽ നിന്ന് വായിക്കുന്നു, സ്കോറുകൾ കണക്കാക്കുന്നു, ഇടക്കാല ഫലങ്ങൾ (intermediate results) തിരികെ എഴുതുന്നു. ഇതിലെ ഗണിതക്രിയകൾ ലളിതമാണ്, എന്നാൽ ഡാറ്റാ കൈമാറ്റം വളരെ കഠിനമാണ്. പവർ, തെർമൽ ബജറ്റുകൾ പരിമിതമായ ആൻഡ്രോയിഡിൽ, ഈ രീതി മെമ്മറി ബസിനെ വല്ലാതെ ബാധിക്കുന്നു. ഒരേ പാലം കടക്കാൻ പ്രോസസ്സർ മൂന്ന് തവണ ടോൾ നൽകുന്നത് പോലെയാണിത്.

നിങ്ങൾ INT8 മോഡലുകൾ ഉപയോഗിക്കുകയും പ്രോംപ്റ്റിന്റെ നീളം കൂടുമ്പോൾ പ്രീഫിൽ ഘട്ടം എന്തുകൊണ്ട് ക്വാഡ്രാറ്റിക് രീതിയിൽ (quadratically) വർദ്ധിക്കുന്നു എന്ന് ആലോചിക്കുകയും ചെയ്യുന്നുണ്ടെങ്കിൽ, ഇതാ ഉത്തരം. വെയ്റ്റുകൾ ചെറുതാണെങ്കിലും, ആക്ടിവേഷൻ ട്രാഫിക് (activation traffic) ഇപ്പോഴും വളരെ കൂടുതലാണ്.

തടസ്സം മെമ്മറിയാണ്, ഗണിതമല്ല

പരിഹാരം മനസ്സിലാക്കാൻ, റൂഫ്‌ലൈൻ (roofline) പരിശോധിക്കുക. Snapdragon 8 Gen 3, Dimensity 9300 തുടങ്ങിയ ചിപ്പുകളിലെ മൊബൈൽ GPU-കളും NPU-കളും അവയുടെ LPDDR5X ഇന്റർഫേസുകൾക്ക് താങ്ങാൻ കഴിയുന്നതിനേക്കാൾ വളരെ ഉയർന്ന തിയററ്റിക്കൽ കമ്പ്യൂട്ട് ത്രൂപുട്ട് (theoretical compute throughput) പ്രകടിപ്പിക്കുന്നു. ഒരു സാധാരണ അറ്റൻഷൻ കേർണലിൽ (attention kernel), ഓരോ ഹെഡും Q, K എന്നിവയുടെ ഗുണനഫലത്തിന്മേൽ softmax കണക്കാക്കുകയും തുടർന്ന് അത് V കൊണ്ട് ഗുണിക്കുകയും ചെയ്യുന്നു. ഓരോ ഇടക്കാല സ്കോർ മാട്രിക്സും ഗ്ലോബൽ മെമ്മറിയിൽ രൂപപ്പെടുന്നു. ഇതിനർത്ഥം നിങ്ങളുടെ പീക്ക് DRAM റീഡുകൾ സീക്വൻസ് നീളത്തിന്റെ വർഗ്ഗത്തിന് (square) അനുസൃതമായി, അതായത് O(n²) എന്ന രീതിയിൽ വർദ്ധിക്കുന്നു എന്നാണ്. 1024 ടോക്കണുകളുള്ള ഒരു പ്രോംപ്റ്റിന്, എക്സിക്യൂഷൻ സമയം നിയന്ത്രിക്കാൻ പാകത്തിൽ മെമ്മറി ട്രാഫിക് ഇതിനകം തന്നെ വളരെ കൂടുതലായിരിക്കും.

ലേറ്റൻസി മറച്ചുവെക്കാൻ കഴിയാത്തതിനാൽ കോറുകൾ വേണ്ടത്ര ഉപയോഗിക്കപ്പെടുന്നില്ല. പൈപ്പ്‌ലൈനുകൾ നിറഞ്ഞുനിൽക്കാൻ ആധുനിക പ്രോസസ്സറുകൾ കാഷെ (cache) ഉപയോഗിക്കുന്നു. ഒരു അൽഗോരിതം നിരന്തരം കാഷെ മിസ്സ് ചെയ്യുകയും DRAM-ൽ നിന്ന് ഡാറ്റ എടുക്കുകയും ചെയ്യുമ്പോൾ, എക്സിക്യൂഷൻ യൂണിറ്റുകൾ വെറുതെ ഇരിക്കേണ്ടി വരുന്നു. കമ്പ്യൂട്ട് ശേഷിയും മെമ്മറി വിതരണവും തമ്മിലുള്ള ഈ ഘടനാപരമായ പൊരുത്തക്കേട് പരിഹരിക്കാൻ ക്വാണ്ടൈസേഷൻ കൊണ്ട് മാത്രം കഴിയില്ല.

ടൈലിംഗ് (Tiling) എങ്ങനെ ബാൻഡ്‌വിഡ്ത്ത് വീണ്ടെടുക്കുന്നു

ഇടക്കാല സ്കോറുകൾ DRAM-ലേക്ക് അയക്കുന്നതിന് പകരം ഓൺ-ചിപ്പ് SRAM-ൽ സൂക്ഷിക്കുന്ന ഒരു ടൈലിംഗ് സ്ട്രാറ്റജി (tiling strategy) ആണ് ഇതിനുള്ള പരിഹാരം. ആൻഡ്രോയിഡ് കമ്പ്യൂട്ട് സ്റ്റാക്കിന് അനുയോജ്യമായ രീതിയിൽ മാറ്റം വരുത്തിയ ഫ്ലാഷ് അറ്റൻഷനെ (Flash Attention) നയിക്കുന്ന അതേ ആശയമാണിത്. മെമ്മറിയിൽ ഒരു പൂർണ്ണമായ n × n സ്കോർ മാട്രിക്സ് രൂപീകരിക്കുന്നതിന് പകരം, നിങ്ങൾ കമ്പ്യൂട്ടേഷനെ L1 കാഷെയ്റ്റിനുള്ളിൽ ഒതുങ്ങുന്ന ചെറിയ ടൈലുകളായി (tiles) വിഭജിക്കുന്നു. നിങ്ങൾ ലോക്കൽ softmax സ്റ്റാറ്റിസ്റ്റിക്സ് കണക്കാക്കുന്നു, റണ്ണിംഗ് മാക്സ് വാല്യൂസും നോർമലൈസേഷൻ സമ്മുകളും ശേഖരിക്കുന്നു, ഒടുവിൽ ഫൈനൽ വെയ്റ്റഡ് ഔട്ട്‌പുട്ടുകൾ മാത്രം മെമ്മറിയിലേക്ക് എഴുതുന്നു.

ഇത് ബാൻഡ്‌വിഡ്ത്ത് കോംപ്ലക്സിറ്റി മാറ്റുന്നു. പീക്ക് DRAM റീഡുകൾ O(n²) എന്നതിൽ നിന്ന് O(n) ആയി കുറയുന്നു, കാരണം പൂർണ്ണമായ സ്കോർ മാട്രിക്സുകൾ മെയിൻ മെമ്മറിയിലൂടെ കൊണ്ടുപോകേണ്ട ആവശ്യം ഇല്ലാതാകുന്നു. പ്രധാനപ്പെട്ട ജോലികളെല്ലാം എക്സിക്യൂഷൻ യൂണിറ്റുകൾക്ക് തൊട്ടടുത്തുള്ള SRAM-നുള്ളിൽ തന്നെ നടക്കുന്നു.

ഒരു വ്യക്തമായ ഉദാഹരണത്തിനായി, 64 എന്ന ടൈൽ സൈസും 128 എന്ന ഹെഡ് ഡൈമെൻഷനും പരിഗണിക്കുക. സ്കോർ ടൈൽ 16 KB സ്ഥലം എടുക്കുന്നു. ഈ വലിപ്പം Snapdragon 8 Gen 3, Dimensity 9300 തുടങ്ങിയ നിലവിലെ ഫ്ലാഗ്ഷിപ്പ് SoCs-ന്റെ L1 കാഷെയ്റ്റിനുള്ളിൽ സുഗമമായി ഒതുങ്ങും. ഗണിതക്രിയകൾ ലോക്കലായി തന്നെ തുടരുന്നു. മെമ്മറി ബസ് സുഗമമായി പ്രവർത്തിക്കുന്നു.

ഇത് ആൻഡ്രോയിഡിൽ നടപ്പിലാക്കുന്നത് എങ്ങനെ

അൽഗോരിതത്തിന്റെ രൂപരേഖ ലളിതമാണ്, എങ്കിലും കൃത്യമായ വിവരങ്ങൾ ശ്രദ്ധിക്കേണ്ടതുണ്ട്.

നിങ്ങളുടെ Query, Key എന്നിവ വിഭജിക്കുക