Leveraging ASIC AI Chips for Homomorphic Encryption

यहाँ HPCA 2026 में प्रस्तुत शोध पत्र "Leveraging ASIC AI Chips for Homomorphic Encryption" का विस्तृत तकनीकी सारांश दिया गया है।

1. समस्या विवरण (Problem Statement)

होमोमॉर्फिक एन्क्रिप्शन (HE) एन्क्रिप्टेड डेटा पर गणना करने में सक्षम बनाता है, जो क्लाउड सेवाओं के लिए मजबूत गोपनीयता प्रदान करता है। हालाँकि, इसका व्यावहारिक उपयोग अत्यधिक कम्प्यूटेशनल ओवरहेड्स (जो अक्सर CPU पर प्लेनटेक्स्ट की तुलना में 1000× धीमा होता है) के कारण बाधित है।

• वर्तमान स्थिति: HE को त्वरित करने के लिए GPU मानक (de-facto standard) हैं, लेकिन वे अभी भी विशेष HE ASICs की तुलना में ऊर्जा दक्षता (लगभग 33× का अंतर) में काफी पीछे हैं।
• अंतराल (The Gap): समर्पित HE ASICs ऊर्जा-कुशल तो हैं, लेकिन उन्हें डिजाइन और निर्मित करना महंगा है, जिससे उनका व्यापक प्रसार सीमित हो जाता है।
• चुनौती: मौजूदा AI एक्सेलेरेटर्स (जैसे Google TPUs) उत्कृष्ट ऊर्जा दक्षता के साथ विशाल कंप्यूट और मेमोरी संसाधन प्रदान करते हैं। हालाँकि, GPU के लिए डिज़ाइन किए गए SoTA HE एल्गोरिदम TPUs पर कुशलतापूर्वक मैप नहीं होते हैं क्योंकि उनके बीच मौलिक आर्किटेक्चरल विसंगतियां हैं:
  1. अंकगणितीय विसंगति (Arithmetic Mismatch): HE उच्च-परिशुद्धता (32-bit) मॉड्यूलर अंकगणित पर निर्भर करता है। GPU इसे अच्छी तरह से संभालते हैं, लेकिन TPUs में 32-bit ऑपरेशन्स के लिए एक लो-थ्रूपुट वेक्टर प्रोसेसिंग यूनिट (VPU) और लो-प्रिसिजन (8-bit) ऑपरेशन्स के लिए एक हाई-थ्रूपुट मैट्रिक्स मल्टीप्लिकेशन यूनिट (MXU) होती है। वर्तमान HE लाइब्रेरीज़ शक्तिशाली MXU को खाली छोड़ देती हैं।
  2. मेमोरी/ग्रैनुलैरिटी विसंगति (Memory/Granularity Mismatch): HE कर्नेल्स (जैसे नंबर थ्योरेटिक ट्रांसफॉर्म - NTT) को बार-बार, सूक्ष्म-स्तर (fine-grained) के डेटा परम्यूटेशन (शफलिंग/ट्रांसपोज़िंग) की आवश्यकता होती है। TPUs मोटे-स्तर (coarse-grained), SIMD-नियंत्रित मेमोरी (बड़े रजिस्टरों) का उपयोग करते हैं, जिससे ये सूक्ष्म-स्तर के ऑपरेशन अत्यंत अक्षम और महंगे हो जाते हैं।

2. कार्यप्रणाली: CROSS फ्रेमवर्क (Methodology: The CROSS Framework)

लेखक CROSS प्रस्तावित करते हैं, जो एक कंपाइलर फ्रेमवर्क है जो HE वर्कलोड को AI ASICs (विशेष रूप से Google TPUs) की मोटे-स्तर (coarse-grained) और उच्च-थ्रूपुट वाली संरचना के अनुरूप व्यवस्थित रूप से रूपांतरित करता है। CROSS दो मुख्य रूपांतरण पेश करता है:

A. बेस-अलाइन्ड ट्रांसफॉर्मेशन (Basis-Aligned Transformation - BAT)

• लक्षक: उच्च-परिशुद्धता (32-bit) मॉड्यूलर अंकगणित को सघन (dense), कम-परिशुद्धता (8-bit) मैट्रिक्स गुणन में बदलना ताकि TPU के MXU का उपयोग किया जा सके।
• तंत्र:
  • मानक HE गुणन उच्च-परिशुद्धता संख्याओं को टुकड़ों में विभाजित करता है, जिसके परिणामस्वरूप कई अनावश्यक शून्य (redundant zeros) के साथ स्पार्स मैट्रिक्स-वेक्टर गुणन होता है (जो कंप्यूट बर्बाद करता है)।
  • BAT पहचानता है कि ट्विडल फैक्टर्स (NTT में) और इवैल्यूएशन कीज़ जैसे पैरामीटर्स स्थिर (static) और ऑफलाइन ज्ञात होते हैं।
  • यह इन स्थिर पैरामीटर्स पर ऑफलाइन $mod\ q$ मॉड्यूलर रिडक्शन लागू करता है, जिससे उच्च-बेसिस टर्म्स को लो-बेसिस टर्म्स में "फोल्ड" (fold) कर दिया जाता है।
  • यह स्पार्स मैट्रिक्स को एक सघन (dense), छोटे मैट्रिक्स में बदल देता है, जिससे अनावश्यक शून्य समाप्त हो जाते हैं और इस ऑपरेशन को MXU पर एक डेंस 8-bit MatMul के रूप में निष्पादित करना संभव हो जाता है।
  • परिणाम: ~2× सैद्धांतिक कंप्यूट और मेमोरी की बर्बादी में कमी; MXU का उपयोग सक्षम करता है (जो TPUs पर VPU की तुलना में ~58× तेज़ है)।

B. मेमोरी-अलाइन्ड ट्रांसफॉर्मेशन (Memory-Aligned Transformation - MAT)

• लक्ष्य: HE एल्गोरिदम द्वारा आवश्यक महंगे रनटाइम डेटा रीऑर्डरिंग (ट्रांसपोज़/शफलिंग) को समाप्त करना।
• तंत्र:
  • रनटाइम पर डेटा परम्यूटेशन करने के बजाय (जो TPUs पर लो रजिस्टर यूटिलाइजेशन और हाई लेटेंसी का कारण बनता है), MAT परम्यूटेशन को परम्यूटेशन मैट्रिक्स द्वारा गुणन के रूप में मानता है।
  • इन परम्यूटेशन मैट्रिसेस को कंपाइलेशन के दौरान ऑफलाइन रूप से पूर्व-ज्ञात पैरामीटर्स (जैसे ट्विडल फैक्टर्स) पर लागू किया जाता है।
  • यह एक "लेआउट-इनवेरिएंट" (layout-invariant) कर्नेल बनाता है जहाँ गणना स्वाभाविक रूप से वांछित क्रम में आउटपुट उत्पन्न करती है, बिना किसी स्पष्ट रनटाइम मेमोरी मूवमेंट के।
  • अनुप्रयोग: 4-स्टेप NTT एल्गोरिदम पर लागू होते हुए, MAT स्पष्ट मैट्रिक्स ट्रांसपोज़ और बिट-रिवर्स शफरिंग स्टेप्स को हटा देता है, जिससे एल्गोरिदम को 3-स्टेप लेआउट-इनवेरिएंट प्रक्रिया में बदल दिया जाता है।

3. मुख्य योगदान (Key Contributions)

1. व्यवस्थित लक्षण वर्णन (Systematic Characterization): उन विशिष्ट आर्किटेक्चरल विसंगतियों (परिशुद्धता और ग्रैनुलैरिटी) की पहचान की जो SoTA GPU HE एल्गोरिदम को AI ASICs पर कुशलतापूर्वक चलने से रोकती हैं।
2. BAT (Basis-Aligned Transformation): उच्च-परिशुद्धता मॉड्यूलर इंटीजर अंकगणित को लो-परिशिजन डेंस मैट्रिक्स गुणन में मैप करने की एक नवीन तकनीक, जो HE के लिए MXU को अनलॉक करती है।
3. MAT (Memory-Aligned Transformation): एक शेड्यूलिंग तकनीक जो डेटा रीऑर्डरिंग को ऑफलाइन पैरामीटर ट्रांसफॉर्मेशन में समाहित करती है, जिससे रनटाइम मेमोरी ओवरहेड समाप्त हो जाता है।
4. प्रथम AI/HE सह-त्वरण (First AI/HE Co-Acceleration): प्रदर्शित किया कि मौजूदा AI ASICs (TPUs) बिना हार्डवेयर संशोधन के HE ऑपरेटर्स को त्वरित कर सकते हैं, जिससे SoTA ऊर्जा दक्षता प्राप्त होती है।

4. प्रयोगात्मक परिणाम (Experimental Results)

लेखकों ने एक वास्तविक Google TPU v6e (सिंगल-होस्ट) पर CROSS का मूल्यांकन किया और इसकी तुलना CPU, GPU (NVIDIA A100, RTX 4090), FPGA और HE ASICs पर SoTA कार्यान्वयन के विरुद्ध की।

• NTT थ्रूपुट:
  • WarpDrive (NVIDIA A100) की तुलना में 1.43× अधिक थ्रूपुट प्राप्त किया।
  • व्यावहारिक रूप से उपलब्ध उपकरणों के लिए SoTA NTT थ्रूपुट रिकॉर्ड को नया रूप दिया।
• ऊर्जा दक्षता (थ्रूपुट प्रति वाट):
  • OpenFHE (CPU) की तुलना में 451× अधिक।
  • WarpDrive (GPU) की तुलना में 7.81× अधिक।
  • FIDESlib (GPU) की तुलना में 1.83× अधिक।
  • विशेषीकृत HE ASICs (BASALISC, CraterLake) में HE-Add दक्षता में उन्हें पीछे छोड़ा और HE-Mult और Rotate में अंतर को काफी कम कर दिया (ASICs के 3× के भीतर)।
• वर्कलोड प्रदर्शन:
  • MNIST इन्फरेंस: Orion (GPU) की तुलना में 10× स्पीडअप, 98% सटीकता के साथ।
  • लॉजिस्टिक्स रिग्रेशन: Cheddar (RTX 4090) की तुलना में 1.06× थ्रूपुट प्रति वाट।
  • बूटस्ट्रैपिंग (Bootstrapping): Cheddar (RTX 4090) की तुलना में 1.5× थ्रूपुट।

5. महत्व और प्रभाव (Significance and Impact)

• प्रतिमान परिवर्तन (Paradigm Shift): यह कार्य सिद्ध करता है कि AI एक्सेलेरेटर्स (जो डीप लर्निंग के लिए डिज़ाइन किए गए हैं) होमोमॉर्फिक एन्क्रिप्शन के लिए सबसे ऊर्जा-कुशल प्लेटफॉर्म हो सकते हैं, जो कमोडिटी हार्डवेयर के लिए प्रदर्शन-प्रति-वाट के मामले में सामान्य-उद्देश्य वाले GPU और विशेष HE ASICs दोनों से बेहतर हैं।
• कोई हार्डवेयर लागत नहीं: समर्पित HE ASICs के विपरीत, CROSS के लिए शून्य हार्डवेयर संशोधन की आवश्यकता है। यह मौजूदा, बड़े पैमाने पर उत्पादित AI इंफ्रास्ट्रक्चर (TPUs) का लाभ उठाकर गोपनीयता-संरक्षित AI को सक्षम बनाता है।
• स्केलेबिलिटी: AI ASICs पर कुशल HE को सक्षम करके, यह दृष्टिकोण बड़े पैमाने पर "प्राइवेट AI" का मार्ग प्रशस्त करता है, जिससे क्लाउड प्रदाता अपने मौजूदा AI क्लस्टर्स का उपयोग करके एन्क्रिप्टेड डेटा पर इन्फरेंस कर सकते हैं।
• भविष्य की दिशाएँ: पेपर में उल्लेख किया गया है कि जबकि CROSS अत्यधिक कुशल है, समर्पित ASICs (3-33×) के बीच का शेष अंतर कम लागत वाले हार्डवेयर-विशिष्ट शफलिंग इंजन और फिक्स्ड मोड्युली सपोर्ट की कमी के कारण है। भविष्य के कार्य विशिष्ट मोडुलस बाधाओं के लिए और अधिक एल्गोरिथमिक परम्यूटेशन एम्बेडिंग और अनुकूलन पर ध्यान केंद्रित कर सकते हैं।

संक्षेप में, CROSS AI ASICs की मोटे-स्तर (coarse-grained) और उच्च-थ्रूपुट प्रकृति और होमोमॉर्फिक एन्क्रिप्शन की सूक्ष्म-स्तर (fine-grained) और उच्च-परिशुद्धता आवश्यकताओं के बीच के अंतर को पाटता है, जिससे ऊर्जा-कुशल गोपनीयता-संरक्षित गणना के लिए एक नया मानक स्थापित होता है।