FlashAttention-2 (UR)

From Systems analysis Wiki
Jump to navigation Jump to search

FlashAttention-2 — یہ ایک بہتر الگورتھم ہے جو بڑے لسانی ماڈلز (LLM) میں attention کے میکانزم کی گنتی کے لیے بنایا گیا ہے۔ یہ الگورتھم Tri Dao اور اسٹینفورڈ یونیورسٹی کے محققین نے تیار کیا اور جولائی 2023 میں پیش کیا گیا[1]۔ اس کا بنیادی مقصد — GPU کے ہارڈویئر وسائل کو زیادہ مؤثر طریقے سے استعمال کرتے ہوئے transformer ماڈلز کی تربیت اور inference کو نمایاں طور پر تیز کرنا ہے، جبکہ معیاری attention میکانزم کے ساتھ مکمل یکسانیت برقرار رکھنا، یعنی درستگی میں کوئی کمی نہیں۔

FlashAttention-2 الگورتھم FlashAttention کا منطقی تسلسل ہے جسے اسی ٹیم نے 2022 میں پیش کیا تھا۔ نئے ورژن نے GPU کی نامکمل لوڈنگ کے مسئلے کو حل کیا جو پہلے ورژن میں موجود تھا، اور پہلے ورژن کے مقابلے میں تقریباً دوگنی رفتار حاصل کی۔

پس منظر: transformers میں attention کا مسئلہ

معیاری self-attention میکانزم transformers میں لمبی متن کی ترتیبوں کے ساتھ کام کرتے وقت ایک رکاوٹ بن جاتا ہے۔ اس کی حسابی پیچیدگی اور میموری کا استعمال ترتیب کی لمبائی (N) کے لحاظ سے مربعاتی طور پر (O(N²)) بڑھتا ہے، جو LLM کے زیادہ سے زیادہ context کی لمبائی اور scalability پر سنگین پابندیاں عائد کرتا ہے[1]۔

اس مسئلے کے حل کے لیے 2022 میں FlashAttention الگورتھم پیش کیا گیا[2]۔ اس کے اہم خیالات:

  • GPU کی میموری درجہ بندی کا لحاظ (IO-awareness): الگورتھم GPU کی سست میموری (HBM) اور چپ پر موجود تیز رفتار جامد میموری (SRAM) کے درمیان پڑھنے/لکھنے کی مہنگی کارروائیوں کو کم سے کم کرتا ہے۔
  • بلاک پروسیسنگ (tiling): گنتی چھوٹے بلاکس (tiles) میں تقسیم کی جاتی ہے جو تیز رفتار SRAM میں پروسیس ہوتے ہیں، جس سے میموری میں مکمل attention میٹرکس کو مادی شکل دینے سے بچا جاتا ہے۔

اس سے میموری کی خطی نمو (O(N)) اور معیاری نفاذ کے مقابلے میں 2–4 گنا تیزی حاصل ہوئی[2]۔ FlashAttention کو وسیع پیمانے پر اپنایا گیا اور اس نے نمایاں طور پر بڑھے ہوئے context والے ماڈلز کی ظہور میں مدد کی، مثلاً 2–4 ہزار tokens (GPT-3) سے 128 ہزار (GPT-4) اور اس سے بھی زیادہ[3]۔ چنانچہ Falcon-40B ماڈل میں FlashAttention کے استعمال سے inference 3 گنا اور GPT-3 کے مقابلے میں مجموعی generation کارکردگی 5 گنا تیز ہوئی[4]۔

FlashAttention-2 کی ترقی اور اہداف

کامیابی کے باوجود، FlashAttention کا پہلا ورژن GPU کے حسابی وسائل کو مکمل طور پر استعمال نہیں کرتا تھا۔ NVIDIA A100 گرافکس کارڈز پر کارکردگی نظریاتی زیادہ سے زیادہ (FLOPs/s) کا صرف 25–40% تک پہنچتی تھی[1]۔ بنیادی وجہ Streaming Multiprocessors کی غیر بہتر لوڈنگ اور مشترکہ میموری کی فاضل کارروائیاں تھیں[5]۔

FlashAttention-2 کا مقصد کام کو زیادہ مؤثر موازی بنانے اور معاون کارروائیوں کو کم سے کم کرکے گنتی کو مزید تیز کرنا تھا۔ الگورتھم کو زیادہ سے زیادہ کارکردگی حاصل کرنے کے لیے NVIDIA CUTLASS 3.x لائبریری کے نچلی سطح کے primitives استعمال کرتے ہوئے مکمل طور پر دوبارہ لکھا گیا[6]۔

تکنیکی فن تعمیر اور کام کے اصول

FlashAttention-2 موازی کاری اور کارکردگی بڑھانے کے لیے تین اہم بہتریاں متعارف کراتا ہے[1]:

1. غیر میٹرکس کارروائیوں کو کم سے کم کرنا

الگورتھم میٹرکس ضرب (non-matmul FLOPs) سے متعلق نہ ہونے والی فلوٹنگ پوائنٹ معاون کارروائیوں کی تعداد کم کرتا ہے۔ چونکہ GPU کے tensor cores خاص طور پر میٹرکس کارروائیوں (GEMM) کے لیے بہتر ہیں اور انہیں 16 گنا تیز انجام دیتے ہیں، یہ تبدیلی GPU کے سب سے زیادہ کارکرد بلاکس کو زیادہ وقت استعمال کرنے کی اجازت دیتی ہے۔

2. بہتر موازی کاری

اصل FlashAttention میں ایک attention "head" پر کام کو موازی نہیں کیا جاتا تھا، جس سے لمبی ترتیبوں اور چھوٹے batch sizes میں بے کاری ہوتی تھی۔ FlashAttention-2 بین بلاک موازی کاری متعارف کراتا ہے: اب ایک attention head کے لیے گنتی GPU کے مختلف Streaming Multiprocessors میں تقسیم کی جاتی ہے، جس سے ان کی لوڈنگ نمایاں طور پر بڑھتی ہے۔

3. بلاک کے اندر کام کی بہتر تقسیم

ایک حسابی بلاک کی سطح پر، مشترکہ میموری (shared memory) کے ذریعے ڈیٹا کے تبادلے کو کم کرنے کے لیے کام کو threads کے گروپوں (warps) کے درمیان دوبارہ تقسیم کیا گیا۔ اس سے Softmax کی نارملائزیشن کے لیے درکار فاضل پڑھنے/لکھنے کی کارروائیوں کی تعداد کم ہوتی ہے۔

کارکردگی اور افادیت

تعمیراتی بہتریوں کی بدولت، FlashAttention-2 کارکردگی میں نمایاں اضافہ ظاہر کرتا ہے:

  • دوگنی رفتار: الگورتھم FlashAttention کے پہلے ورژن کے مقابلے میں تقریباً 2 گنا تیز کام کرتا ہے[1]۔
  • GPU کا بھرپور استعمال: NVIDIA A100 GPU پر نظریاتی زیادہ سے زیادہ throughput (TFLOPs) کا 50–73% حاصل ہوتا ہے، جو بہتر میٹرکس ضرب (GEMM) کارروائیوں کی کارکردگی کے قریب ہے[1]۔
  • ریکارڈ گنتی کی رفتار:
    • A100 GPU پر GPT قسم کے ماڈل کے end-to-end تربیتی چکر میں 225 TFLOP/s تک کی رفتار حاصل ہوتی ہے، جو حسابی بلاکس کے 72% استعمال کے مساوی ہے۔ موازنے کے لیے، انہی حالات میں معیاری attention GPU کو 100 TFLOP/s سے کم پر لوڈ کرتا تھا[7]۔
    • H100 GPU پر کارکردگی 335 TFLOP/s تک پہنچتی ہے[7]۔

کارکردگی میں اس اضافے سے، مثلاً 16k tokens کی context window والے ماڈل کو اتنے ہی وقت میں تربیت دینا ممکن ہو جاتا ہے جتنا پہلے 8k tokens کی window کے لیے درکار تھا[5]۔ اہم بات یہ ہے کہ الگورتھم درست اور deterministic رہتا ہے، اس لیے اس کے استعمال سے ماڈل کی پیش گوئیوں کے معیار پر کوئی اثر نہیں پڑتا[8]۔

LLM ماحولیاتی نظام میں استعمال اور انضمام

FlashAttention-2 جلد ہی LLM کے ماحولیاتی نظام میں ایک معیاری آلہ بن گیا۔ یہ بہت سے مقبول frameworks اور لائبریریوں میں شامل کیا گیا ہے:

  • PyTorch: مقامی سپورٹ۔
  • Hugging Face Transformers: ماڈل لوڈ کرتے وقت `attn_implementation="flash_attention_2"` پیرامیٹر سے سپورٹ فعال ہوتی ہے[9]۔ درجنوں آرکیٹیکچرز (GPT, Llama, Falcon, BERT وغیرہ) کے ساتھ مطابقت[10]۔
  • TensorRT-LLM، xFormers اور Triton: الگورتھم ان پلیٹ فارمز کے لیے نافذ کیا گیا ہے، جو وسیع استعمال کو ممکن بناتا ہے[7]۔

انضمام FlashAttention-2 کو دیگر بہتری کے طریقوں جیسے quantization (GPTQ, QLoRA) اور مؤثر fine-tuning (PEFT) کے ساتھ آسانی سے ملانے کی اجازت دیتا ہے[9]۔

بعد کے ورژنوں سے موازنہ

FlashAttention-3

attention کی بہتری کے شعبے میں تحقیق جاری ہے۔ جولائی 2024 میں Tri Dao نے FlashAttention-3 پیش کیا جو NVIDIA Hopper (H100/H200) GPU آرکیٹیکچر کی صلاحیتوں کو بروئے کار لانے پر مرکوز ہے۔ اہم نئی خصوصیات[3]:

  • FP8 سپورٹ: مزید تیزی کے لیے 8-bit floating point گنتی استعمال کرتا ہے۔
  • غیر ہم وقت کارروائیاں: GPU کی غیر ہم وقت صلاحیتوں کو زیادہ مؤثر طریقے سے استعمال کرتا ہے۔

FlashAttention-3 H100 GPU پر FlashAttention-2 کے مقابلے میں 1.5–2 گنا تیزی فراہم کرتا ہے، 740 TFLOP/s (نظریاتی زیادہ سے زیادہ کا 75%) تک کی کارکردگی حاصل کرتا ہے[11]۔

ادبیات

  • 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.
  • Shah, J. et al. (2024). FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision. arXiv:2407.08608.
  • 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. OpenReview: rog0J435OO.
  • Abbott, V.; Zardini, G. (2025). FlashAttention on a Napkin: A Diagrammatic Approach to Deep Learning IO-Awareness. arXiv:2412.03317.

حواشی

  1. ↑ 1.0 1.1 1.2 1.3 1.4 1.5 Дао, Три. «FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning». arXiv:2307.08691 [cs.LG], 17 июля 2023 г. [1]
  2. ↑ 2.0 2.1 «Optimizing LLMs for Speed and Memory». Hugging Face Documentation. [2]
  3. ↑ 3.0 3.1 Дао, Три. «FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision». Tri Dao's Blog. [3]
  4. ↑ «FlashAttention vs FlashAttention-2 - an Analysis». E2E Networks Blog. [4]
  5. ↑ 5.0 5.1 «FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning». OpenReview. [5]
  6. ↑ «FlashAttention-2». Hazy Research, Stanford University. [6]
  7. ↑ 7.0 7.1 7.2 Дао, Три. «FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning» (PDF). arXiv:2307.08691. [7]
  8. ↑ Рашка, Себастьян. «Llama 2 and FlashAttention 2». Ahead of AI Magazine. [8]
  9. ↑ 9.0 9.1 Белькада, Юнес. «Faster and more memory efficient models with Flash Attention 2!». LinkedIn. [9]
  10. ↑ «GPU inference». Hugging Face Documentation. [10]
  11. ↑ Дао, Три, и др. «FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision». arXiv:2407.08608 [cs.LG], 11 июля 2024 г. [11]
Retrieved from "https://systems-analysis.info/int/index.php?title=FlashAttention-2_(UR)&oldid=2215"