সেলফ-অ্যাটেনশন, CUDA কার্নেল ও FlashAttention

আধুনিক Transformer বিশাল টোকেন ক্রমকে হার্ডওয়্যারের সর্বোচ্চ গতিতে প্রক্রিয়া করতে Self-Attention-এর ম্যাট্রিক্স গণনা, Triton-এর মতো উচ্চস্তরের GPU কম্পাইলার এবং FlashAttention-এর মতো মেমরি-সচেতন CUDA কার্নেলের ওপর নির্ভর করে।

1. Vision Transformer-এ Self-Attention গাণিতিকভাবে কীভাবে কাজ করে

Vision Transformer (ViT)-এ একটি ছবিকে (যেমন 224×224 পিক্সেল) 16×16 প্যাচের গ্রিডে কাটা হয়, যাতে N = 196টি সমতল প্যাচ তৈরি হয়। প্রতিটি প্যাচকে রৈখিকভাবে D মাত্রার (যেমন 768) ভেক্টরে প্রক্ষেপণ করে অবস্থানের এমবেডিংয়ের সঙ্গে মেলানো হয়, ফলে ইনপুট ম্যাট্রিক্স X ∈ R^(N × D) তৈরি হয়।

• Query, Key ও Value তৈরি। মডেল প্যাচ ম্যাট্রিক্স X-কে তিনটি শেখা ওয়েট ম্যাট্রিক্স (W_Q, W_K, W_V ∈ R^(D × d_k)) দিয়ে গুণ করে ছবির প্রতিটি প্যাচের জন্য তিনটি আলাদা উপস্থাপন তৈরি করে:

Q=XW_(Q), K=XW_(K), V=XW_(V)

• Query (Q): ছবির একটি নির্দিষ্ট প্যাচ কী খুঁজছে তা প্রকাশ করে (যেমন, "আমি পশমের টেক্সচার, কাছাকাছি চোখ বা কান খুঁজছি")।
• Key (K): একটি প্যাচে কী তথ্য আছে তা জানায় (যেমন, "আমার মধ্যে একটি তীক্ষ্ণ বাঁকানো প্রান্ত আছে")।
• Value (V): প্রকৃত দৃশ্যগত বৈশিষ্ট্যের ডেটা ধরে রাখে, যা মিল পেলে সামনে পাঠানো হবে।

• অ্যাটেনশন স্কোর গণনা। প্রতিটি প্যাচের সঙ্গে অন্য প্রতিটি প্যাচের সম্পর্ক কতটা শক্তিশালী তা মাপতে মডেল সব Query ও Key-এর মধ্যে ডট প্রোডাক্ট গণনা করে:

Scores=(QK^(T))/(sqrt(d_(k)))

sqrt(d_(k)) দিয়ে ভাগ করলে বড় ডট প্রোডাক্ট পরবর্তী softmax ফাংশনকে অতি ক্ষুদ্র গ্রেডিয়েন্টের অঞ্চলে ঠেলে দিতে পারে না। এতে N × N অ্যাটেনশন ম্যাট্রিক্স তৈরি হয়, যেখানে (i, j) ঘরটি প্যাচ i-এর প্যাচ j-তে কতটা মনোযোগ দেওয়া উচিত তার স্কোর দেয়।

• Softmax ওজনায়ন ও একত্রীকরণ। সারিভিত্তিক Softmax প্রয়োগ করলে কাঁচা স্কোর এমন সম্ভাবনায় রূপান্তরিত হয় যাদের যোগফল 1। এই সম্ভাবনাগুলোকে Value ম্যাট্রিক্স V দিয়ে গুণ করলে প্রতিটি প্যাচের জন্য মিশ্রিত উপস্থাপন তৈরি হয়:

Attention(Q,K,V)=softmax((QK^(T))/(sqrt(d_(k))))V

Multi-Head Self-Attention-এ ঠিক এই প্রক্রিয়াই একাধিক হেডে (যেমন 12টি হেড) সমান্তরালে চলে, যাতে বিভিন্ন হেড একই সময়ে বিভিন্ন স্থানিক দূরত্ব, রং বা বস্তুর অংশে মনোযোগ দিতে পারে।

2. হাতে-কলমে কোড তুলনা: সরাসরি CUDA C++ বনাম OpenAI Triton

GPU প্রোগ্রামিং কীভাবে বিবর্তিত হয়েছে তা দেখতে Nটি উপাদানের জন্য একটি সহজ উপাদানভিত্তিক ভেক্টর যোগ (C = A + B) বিবেচনা করুন।

সরাসরি CUDA C++ (থ্রেড-স্তরের ব্যবস্থাপনা)

CUDA C++-এ আপনি একটি স্কেলার থ্রেডের দৃষ্টিকোণ থেকে প্রোগ্রাম করেন। থ্রেড ID গণনা, অ্যারের সীমানা পরীক্ষা এবং হোস্ট-সাইড লঞ্চ কনফিগারেশনের প্রয়োজনীয় পুনরাবৃত্তিমূলক কোড নিজে লিখতে হয়।

#include <cuda_runtime.h>

// Device kernel: executed by thousands of individual threads
__global__ void vector_add_cuda(const float* A, const float* B, float* C, int N) {
    // Manually compute global thread index from block and thread IDs
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    
    // Guard against out-of-bounds memory access
    if (idx < N) {
        C[idx] = A[idx] + B[idx];
    }
}

// Host launch setup
void launch_cuda(const float* d_A, const float* d_B, float* d_C, int N) {
    int threads_per_block = 256;
    int blocks_per_grid = (N + threads_per_block - 1) / threads_per_block;
    vector_add_cuda<<<blocks_per_grid, threads_per_block>>>(d_A, d_B, d_C, N);
}

OpenAI Triton (ব্লক-স্তরের Python প্রোগ্রামিং)

Triton-এ সরাসরি Python-এ ডেটার একটি ব্লক (টাইল)-এর দৃষ্টিকোণ থেকে প্রোগ্রাম করেন। Triton-এর কম্পাইলার ভেতরে স্বয়ংক্রিয়ভাবে থ্রেড সমলয়করণ, মেমরি কোয়ালেসিং ও ভেক্টরায়ন সামলায়।

import triton
import triton.language as tl

@triton.jit
def vector_add_triton(A_ptr, B_ptr, C_ptr, N, BLOCK_SIZE: tl.constexpr):
    # Identify which block of data this program instance handles
    pid = tl.program_id(axis=0)
    
    # Create a vector of offsets for the entire block at once
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < N
    
    # Load chunks of memory into fast SRAM, add them, and store back
    a = tl.load(A_ptr + offsets, mask=mask)
    b = tl.load(B_ptr + offsets, mask=mask)
    tl.store(C_ptr + offsets, a + b, mask=mask)

# Launch directly from PyTorch tensors
def launch_triton(A, B, C):
    N = A.numel()
    grid = lambda meta: (triton.cdiv(N, meta['BLOCK_SIZE']),)
    vector_add_triton[grid](A, B, C, N, BLOCK_SIZE=1024)

3. FlashAttention: Transformer-এর মেমরি বাধার সমাধান

উচ্চ-রেজোলিউশনের ছবি বা দীর্ঘ টেক্সট ক্রম প্রক্রিয়া করার সময় সাধারণ Self-Attention কঠিন সীমায় পৌঁছায়, কারণ অ্যাটেনশন ম্যাট্রিক্স S = QK^T ক্রমের দৈর্ঘ্য N-এর সঙ্গে বর্গানুপাতে (N × N) বাড়ে।

• মেমরি ব্যান্ডউইডথের বাধা (HBM বনাম SRAM)। আধুনিক NVIDIA GPU-তে মেমরির প্রধান দুটি স্তর রয়েছে:

• HBM (High Bandwidth Memory / VRAM): বড় (যেমন 80 GB), কিন্তু পড়া ও লেখায় তুলনামূলক ধীর।
• SRAM (চিপের ভেতরের Shared Memory): অতি ছোট (সব Streaming Multiprocessor মিলিয়ে মোট ~20 MB), কিন্তু HBM-এর চেয়ে প্রায় 10× থেকে 20× দ্রুত।

সাধারণ অ্যাটেনশনে GPU প্রথমে S = QK^T গণনা করে, বিশাল N × N ম্যাট্রিক্স ধীর HBM-এ লেখে, Softmax প্রয়োগ করতে তা আবার পড়ে, সম্ভাবনাগুলো আবার HBM-এ লেখে এবং V দিয়ে গুণ করতে সেগুলো আবার পড়ে। Tensor Cores তাদের অধিকাংশ সময় মেমরি স্থানান্তরের অপেক্ষায় নিষ্ক্রিয় থাকে।

• টাইলিং (N × N ম্যাট্রিক্স কখনো পুরো তৈরি করে না রাখা)। Tri Dao ও সহকর্মীদের (2022) চালু করা FlashAttention কার্নেল ফিউশন দিয়ে পুরো অ্যাটেনশন গণনা দ্রুত অন-চিপ SRAM-এর ভেতরে রাখে। এটি Q, K ও V-কে SRAM-এ ধরে এমন ছোট ব্লকে ভাগ করে, টাইল ধরে ধরে অ্যাটেনশন গণনা করে এবং শুধু চূড়ান্ত আউটপুট ম্যাট্রিক্স HBM-এ ফেরত লেখে। বিশাল N × N মধ্যবর্তী ম্যাট্রিক্স কখনো VRAM-এ সংরক্ষণ করা হয় না।

• অনলাইন Softmax কৌশল। সাধারণভাবে কোনো একটি উপাদান স্বাভাবিকীকরণের আগেই Softmax-এর পুরো সারির সর্বোচ্চ মান ও সূচকীয় মানগুলোর যোগফল জানা লাগে। FlashAttention একটি অনলাইন Softmax অ্যালগরিদম ব্যবহার করে, যা ব্লক ধরে চলমান সর্বোচ্চ মান ও স্বাভাবিকীকরণের গুণকগুলো অনুসরণ করে এবং প্রতিটি নতুন টাইল প্রক্রিয়া হলে আংশিক আউটপুট গাণিতিকভাবে পুনরায় স্কেল করে।

• ব্যাকপ্রপাগেশনের সময় পুনর্গণনা। FlashAttention ফরোয়ার্ড পাসে N × N অ্যাটেনশন ম্যাট্রিক্স সংরক্ষণ করে না, তাই প্রশিক্ষণের সময় একটি বুদ্ধিদীপ্ত সমঝোতা ব্যবহার করে: এটি শুধু ছোট স্বাভাবিকীকরণ পরিসংখ্যান HBM-এ রাখে এবং ব্যাকপ্রপাগেশনের সময় দ্রুত SRAM-এর ভেতরে প্রয়োজনমতো অ্যাটেনশন টাইলগুলো আবার গণনা করে। SRAM-এ বাড়তি গণনা করা HBM থেকে বড় ম্যাট্রিক্স পড়ার চেয়ে অনেক দ্রুত প্রমাণিত হয়, যা প্রকৃত অতিবাহিত সময়ে 2× থেকে 4× গতি বৃদ্ধি দেয় এবং মেমরির ব্যবহার বর্গানুপাতিক O(N²) থেকে রৈখিক O(N)-এ নামিয়ে আনে।