FlashAttention-3 (UR)

FlashAttention-3 — یہ ٹرانسفارمر نیورل نیٹ ورکس میں attention کے طریقہ کار کو بہتر بنانے کا ایک algorithm ہے، جسے NVIDIA Hopper (H100) GPU آرکیٹیکچر کی ہارڈ ویئر صلاحیتوں کو زیادہ سے زیادہ استعمال کرنے کے لیے تیار کیا گیا ہے۔ یہ algorithm سنہ 2024 میں Colfax Research، Meta، NVIDIA، Georgia Tech، پرنسٹن یونیورسٹی اور Together AI کے محققین کے ایک گروہ نے پیش کیا۔ یہ تحقیقی کام NeurIPS 2024 کانفرنس میں قبول کیا گیا اور spotlight کے طور پر نمایاں کیا گیا۔

FlashAttention-3، الگورتھم کے اس خاندان کی تیسری نسل ہے جو FlashAttention (2022) اور FlashAttention-2 (2023) کے بعد آئی ہے۔ اس کا بنیادی مقصد بڑے لسانی ماڈلز (LLM) کی تربیت اور inference کو نمایاں طور پر تیز کرنا ہے، جبکہ حسابی درستگی کو برقرار رکھنا ہے۔

تعارف اور پس منظر

attention کے طریقہ کار کا مسئلہ

ٹرانسفارمرز کا ایک اہم جزو self-attention کا طریقہ کار ہے، تاہم اس کی حسابی پیچیدگی اور میموری کی کھپت ان پٹ سیکوئنس (n) کی لمبائی بڑھنے کے ساتھ مربع (O(n²)) کے تناسب سے بڑھتی ہے۔ یہ ایک سنگین بوتل نیک پیدا کرتا ہے، کیونکہ جدید GPU تیز رفتار میٹرکس ضرب کے لیے بہترین ہیں، لیکن ایکسپونینشل فنکشنز (مثلاً Softmax میں) کا حساب کئی گنا سست ہوتا ہے۔ اس کے علاوہ، سادہ طریقے سے لاگو کرنے پر GPU میموری میں attention کا ایک بڑا درمیانی tensor محفوظ کرنا پڑتا ہے، جو ماڈلز کی توسیع پذیری کو محدود کرتا ہے۔

FlashAttention اور FlashAttention-2

اس مسئلے کے حل کے لیے سنہ 2022 میں FlashAttention پیش کیا گیا، جس نے دو تکنیکوں کے ذریعے سست عالمی میموری (HBM) تک رسائی کی تعداد کم کی:

• بلاک پر مبنی پروسیسنگ (tiling): حسابات کو بلاکس (tiles) میں تقسیم کیا جاتا ہے، جنہیں تیز رفتار on-chip میموری (SRAM) میں پروسیس کیا جاتا ہے۔
• آپریشنز کا ادغام: تمام آپریشنز (میٹرکس ضرب، Softmax) ایک ہی GPU kernel میں انجام پاتے ہیں اور درمیانی نتائج عالمی میموری میں محفوظ نہیں کیے جاتے۔

اس سے میموری کی پیچیدگی مربع سے لکیری ہو گئی اور حسابات 2 سے 4 گنا تیز ہو گئے۔

سنہ 2023 میں بہتر ورژن — FlashAttention-2 — پیش کیا گیا، جس نے حسابات کی متوازی کاری کو بہتر بنایا۔ NVIDIA Ampere (A100) آرکیٹیکچر کے GPU پر اس نے زیادہ سے زیادہ نظری کارکردگی کا ~70% حاصل کیا۔ تاہم، نئی NVIDIA Hopper (H100) آرکیٹیکچر پر اس کی کارکردگی کافی کم — تقریباً 35% — رہی۔ اس کی وجہ یہ تھی کہ الگورتھم Hopper کی نئی ہارڈ ویئر صلاحیتوں کو استعمال نہیں کرتا تھا، جس نے FlashAttention-3 کی تخلیق کی راہ ہموار کی۔

GPU Hopper (H100) کی نئی ہارڈ ویئر صلاحیتیں

NVIDIA Hopper آرکیٹیکچر نے کئی نئی خصوصیات فراہم کیں جنہیں FlashAttention-3 زیادہ سے زیادہ کارکردگی حاصل کرنے کے لیے استعمال کرتا ہے:

• WGMMA (Warpgroup Matrix Multiply-Accumulate): tensor cores کے لیے ہدایات کی ایک نئی قسم، جو Ampere آرکیٹیکچر کے مقابلے میں تقریباً دوگنی کارکردگی کے ساتھ میٹرکس ضرب انجام دیتی ہے۔
• TMA (Tensor Memory Accelerator): ایک ہارڈ ویئر ماڈیول جو عالمی (HBM) اور مشترکہ (shared memory) میموری کے درمیان ڈیٹا منتقلی کو تیز کرتا ہے۔ TMA خود بخود ایڈریس حسابات کرتا ہے، جس سے computational kernels پر بوجھ کم ہوتا ہے۔
• FP8 فارمیٹ: 8-bit فلوٹنگ پوائنٹ ڈیٹا فارمیٹ کے لیے ہارڈ ویئر سپورٹ، جو FP16 کے مقابلے میں نظری کارکردگی کو دوگنا کرتا ہے، لیکن محدود dynamic range کی وجہ سے درستگی کی کمی کا خطرہ رکھتا ہے۔

FlashAttention-3 کی تکنیکی اختراعات

یہ الگورتھم تین اہم بہتری کے طریقے لاگو کرتا ہے، جو خاص طور پر Hopper آرکیٹیکچر کے لیے تیار کیے گئے ہیں:

1. غیر متزامن عمل اور warps کی تخصص

FlashAttention-3 warp-specialization کے اصول کو استعمال کرتا ہے، جس کے تحت GPU پر دھاگوں (warps) کے مختلف گروہ مختلف کاموں میں مہارت رکھتے ہیں:

• Producer warps: TMA کے ذریعے عالمی میموری سے ڈیٹا لوڈ کرتے ہیں۔
• Consumer warps: tensor cores پر میٹرکس ضرب انجام دیتے ہیں۔

Hopper کی ہارڈ ویئر غیر تزامن کاری کی بدولت، یہ آپریشنز وقت میں ایک دوسرے پر اوورلیپ ہوتے ہیں۔ جب warps کا ایک گروہ حسابات انجام دے رہا ہوتا ہے، دوسرا گروہ بیک وقت اگلے بلاک کے لیے ڈیٹا لوڈ کرتا ہے۔ یہ pipeline طریقہ کار، جو ping-pong scheduling کے اصول پر منظم ہے، سست آپریشنز (مثلاً Softmax) کی تاخیر کو چھپانے اور GPU کے تمام functional modules کو پوری طرح سرگرم رکھنے کی اجازت دیتا ہے۔

2. میموری آپریشنز کو کم سے کم کرنا

الگورتھم پچھلے ورژنز سے tiling کا طریقہ برقرار رکھتا ہے، لیکن TMA کو فعال طور پر اگلے ڈیٹا بلاکس کو موجودہ حسابات کے متوازی طور پر غیر متزامن لوڈ کرنے کے لیے استعمال کرتا ہے۔ سست HBM سے تیز SRAM میں ڈیٹا کی منتقلی عملاً بنیادی حسابات کے سائے میں انجام پاتی ہے، جس کی وجہ سے GPU ڈیٹا کے انتظار میں کم بیکار رہتا ہے۔

3. کم درستگی (FP8) اور quantization کی خطا میں کمی

FP8 پر منتقلی رفتار دوگنی کرتی ہے، لیکن quantization کی وجہ سے درستگی میں نمایاں کمی آ سکتی ہے۔ اس سے نمٹنے کے لیے ڈویلپرز نے incoherent processing کا طریقہ متعارف کرایا۔ اس کا جوہر درج ذیل ہے:

1. attention کا حساب لگانے سے پہلے، خصوصیت کے vectors (queries Q اور keys K) کو ایک بے ترتیب orthogonal matrix (مثلاً Hadamard matrix) سے ضرب دیا جاتا ہے۔
2. یہ تبدیلی غیر معمولی طور پر بڑی قدروں (outliers) کو تمام coordinates میں پھیلا دیتی ہے، ان کی تقسیم کو ہموار کرتے ہوئے۔
3. اس کے بعد FP8 میں quantization انجام پاتی ہے، جو اب کم خطا کے ساتھ ہوتی ہے۔
4. چونکہ تبدیلی orthogonal ہے، یہ attention کے حتمی نتیجے (QKᵀ) کو مسخ نہیں کرتی، کیونکہ ضرب کے وقت matrix کا اثر خود بخود ختم ہو جاتا ہے۔

اس تکنیک نے FP8 میں attention کی حسابی خطا کو بغیر کسی تبدیلی کے معیاری FP8 کے استعمال کے مقابلے میں تقریباً 2.6 گنا کم کر دیا۔

کارکردگی اور اہمیت

ان تکنیکوں کے اطلاق نے FlashAttention-3 کو H100 GPU پر پچھلے ورژنز پر نمایاں برتری دلائی:

• FlashAttention-2 کے مقابلے میں 1.5 سے 2 گنا تیز رفتاری۔
• GPU کا بلند استعمال: H100 کی نظری زیادہ سے زیادہ کارکردگی کا ~75–85% حاصل ہوتا ہے۔
• throughput:
  • نصف درستگی (FP16/BF16) کے لیے 740–840 TFLOPS تک۔
  • 8-bit درستگی (FP8) استعمال کرنے پر 1.2–1.3 PFLOPS (petaflops) تک۔

FlashAttention-3 کی بلند کارکردگی براہ راست LLM کی ترقی اور استعمال پر اثر ڈالتی ہے:

• تربیت کے وقت میں کمی: attention کو 75 سے 100 فیصد تیز کرنا ماڈلز کی تربیت کے وقت کو نمایاں طور پر کم کرتا ہے، جو ہفتوں یا مہینوں پر محیط ہو سکتی ہے۔
• context window میں اضافہ: ماڈلز زیادہ لمبی sequences (لاکھوں tokens) کو مؤثر طریقے سے پروسیس کر سکتے ہیں، جو بڑے دستاویزات یا کوڈ کے تجزیے کے لیے اہم ہے۔
• وسائل کا عقلمندانہ استعمال: کم GPU پر وہی کارکردگی حاصل کرنا یا اسی ہارڈ ویئر پر زیادہ رفتار پانا ممکن ہوتا ہے، جس سے ماڈلز کی deployment کی لاگت کم ہوتی ہے۔

دستیابی اور انضمام

مصنفین نے FlashAttention-3 کا سورس کوڈ GitHub پر کھلے لائسنس کے تحت شائع کیا ہے۔ اس کے PyTorch اور Hugging Face Transformers جیسے معروف deep learning frameworks میں انضمام کی توقع ہے، جو اس ٹیکنالوجی کو ڈویلپرز اور محققین کی ایک وسیع تعداد تک قابل رسائی بنائے گا۔ پچھلے ورژن پہلے ہی صنعت میں de facto معیار بن چکے ہیں، اور FlashAttention-3 غالباً اس رجحان کو آگے بڑھائے گا۔

روابط

• GitHub پر FlashAttention کا سرکاری repository
• FlashAttention-3 کے اعلان کے ساتھ Together AI کا بلاگ

ادبیات

• Shah, J. et al. (2024). FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision. arXiv:2407.08608.
• Dao, T. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. arXiv:2307.08691.
• Dao, T. et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. arXiv:2205.14135.
• Kwon, W. et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. arXiv:2309.06180.
• Ye, Z. et al. (2025). FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving. arXiv:2501.01005.
• Chen, Y. et al. (2023). FlashDecoding++: Faster Large Language Model Inference on GPUs. arXiv:2311.01282.
• Liu, Y. et al. (2024). FastAttention: Extending FlashAttention-2 to NPUs and Low-Resource GPUs. OpenReview: 76NYyOrnfk.
• Dege, P. et al. (2025). FlashMLA-ETAP: Efficient Transpose Attention Pipeline for Accelerating MLA Inference on NVIDIA H20 GPUs. arXiv:2506.01969.
• Wang, G. et al. (2024). FlashMask: Efficient and Rich Mask Extension of FlashAttention. arXiv:2410.01359.
• Abbott, V.; Zardini, G. (2025). FlashAttention on a Napkin: A Diagrammatic Approach to Deep Learning IO-Awareness. arXiv:2412.03317.