ট্রান্সফরমার উদ্ভাবনের পর পাঁচ বছর ধরে সেলফ-অ্যাটেনশনকে নিছক একটি গাণিতিক সমীকরণ মনে করা হতো: $text{Softmax}(mathbf{Q}mathbf{K}^T / sqrt{d})mathbf{V}$। কিন্তু বাস্তবে দীর্ঘ টেক্সটের ক্ষেত্রে অ্যাটেনশন ধীরগতির হয়ে পড়েছিল গাণিতিক জটিলতার ($O(N^2)$ FLOPs) কারণে নয়, বরং মেমরি ইনপুট/আউটপুট (IO) ট্রাফিকের চাপে। ২০২২ সালে ট্রাই দাও এবং ক্রিস্টোফার রে উদ্ভাবন করেন FlashAttention, যা হার্ডওয়্যারের অভ্যন্তরীণ মেমরি বিন্যাসকে সরাসরি বিবেচনায় নিয়ে অ্যাটেনশন গণনায় কোনো প্রকার আনুমানিক ভুল ছাড়াই ২ থেকে ৪ গুণ গতি এবং মেমরি খরচ ১০ গুণ পর্যন্ত কমিয়ে আনে।
১. জিপিইউর অভ্যন্তরীণ মেমরি স্তর: SRAM বনাম HBM
একটি আধুনিক জিপিইউতে (যেমন NVIDIA A100 বা H100) মূলত দুটি মেমরি স্তর থাকে:
- হাই ব্যান্ডউইথ মেমরি (HBM / DRAM): বিশাল ধারনক্ষমতা (৮০ জিবি), কিন্তু তুলনামূলকভাবে ধীর গতি (১.৫–৩.৩৫ টেরাবাইট/সেকেন্ড) এবং উচ্চ লেটেন্সি।
- স্ট্যাটিক র্যাম (SRAM / Shared Memory): অত্যন্ত ক্ষিপ্র গতি (প্রায় ১৯ টেরাবাইট/সেকেন্ড), কিন্তু আকার খুবই ক্ষুদ্র (প্রতি প্রসেসরে মাত্র ১৯২ কিলোবাইট)।
চিরায়ত পাইটর্চে অ্যাটেনশন হিসাব করার সময় প্রতি ধাপে মধ্যবর্তী বিশাল টেনসরগুলোকে ধীরগতির HBM-এ লিখতে ও পড়তে হয়:
- HBM থেকে $mathbf{Q}, mathbf{K}$ এনে গুণ করে $mathbf{S} = mathbf{Q}mathbf{K}^T in mathbb{R}^{N times N}$ তৈরি করা হয় $rightarrow$ আবার HBM-এ লেখা হয়।
- HBM থেকে $mathbf{S}$ পড়ে সফটম্যাক্স করে $mathbf{P}$ বের করা হয় $rightarrow$ আবার HBM-এ লেখা হয়।
- HBM থেকে $mathbf{P}$ ও $mathbf{V}$ পড়ে গুণ করে চূড়ান্ত আউটপুট $mathbf{O}$ বের করা হয়।
সিকোয়েন্স দৈর্ঘ্য ৩২,০০০ হলে এই $N times N$ ম্যাট্রিক্সের আকার দাঁড়ায় প্রতি ব্যাচ আইটেমে ২ গিগাবাইট! ফলে জিপিইউর কম্পিউট কোরগুলো অলস বসে থেকে কেবল ডেটা ট্রাফিকের শিকার হয়।
২. অনলাইন সফটম্যাক্স (Online Softmax) কৌশল
অ্যাটেনশনকে ক্ষুদ্র ব্লকে ভাগ করার প্রধান অন্তরায় ছিল সফটম্যাক্সের হর: $sum_{j=1}^N exp(S_{ij})$। পুরো সারির সমস্ত সংখ্যা না জানা পর্যন্ত সাধারণ সফটম্যাক্স হিসাব করা যায় না।
FlashAttention এই সমস্যার সমাধান করে অনলাইন সফটম্যাক্স সূত্রের সাহায্যে। ধরা যাক দুটি ব্লকের মান $x^{(1)}$ ও $x^{(2)}$, যাদের লোকাল সর্বোচ্চ মান $m_1$ ও $m_2$। সম্মিলিত সর্বোচ্চ মান হবে:
$m_{text{new}} = max(m_1, m_2)$
মধ্যবর্তী এক্সপোনেনশিয়াল যোগফলকে স্কেলিং গুণক দিয়ে সাথে সাথে আপডেট করা যায়:
$l_{text{new}} = l_1 e^{m_1 – m_{text{new}}} + l_2 e^{m_2 – m_{text{new}}}$
একইভাবে সঞ্চিত আউটপুট ভেক্টর $mathbf{O}$-কেও ধাপে ধাপে আপডেট করা হয়:
$mathbf{O}_{text{new}} = mathbf{O}_1 left(frac{l_1 e^{m_1 – m_{text{new}}}}{l_{text{new}}}right) + mathbf{O}_2 left(frac{e^{m_2 – m_{text{new}}}}{l_{text{new}}}right)$
৩. টাইলিং অ্যালগরিদম ও ফলাফল
অনলাইন সফটম্যাক্সের কল্যাণে FlashAttention ইনপুট ম্যাট্রিক্সগুলোকে ক্ষুদ্র ব্লকে বিভক্ত করে যা সরাসরি দ্রুতগতির SRAM-এর ভেতরে এঁটে যায়। $N times N$ ম্যাট্রিক্সটি কখনোই ধীরগতির HBM মেমরিতে লেখা হয় না। মেমরি ট্রাফিক $O(N^2)$ থেকে কমে $O(Nd)$ তে রূপান্তরিত হয়, যার ফলে দীর্ঘ প্রম্পট ও বৃহৎ এলএলএম দ্রুততম গতিতে প্রশিক্ষণ ও ইনফারেন্স পরিচালনা করতে পারে।
