FlashAttention-2 (HE)

From Systems analysis Wiki
Jump to navigation Jump to search

FlashAttention-2 — הוא אלגוריתם משופר המיועד לחישוב מנגנון הattention במודלי שפה גדולים (LLM). האלגוריתם פותח על ידי טרי דאו (Tri Dao) וחוקרים מאוניברסיטת סטנפורד והוצג ביולי 2023[1]. מטרתו המרכזית היא לזרז באופן משמעותי את האימון וה-inference של מודלי transformer באמצעות שימוש יעיל יותר במשאבי חומרת ה-GPU, תוך שמירה על זהות מלאה של החישובים עם מנגנון הattention הסטנדרטי, כלומר ללא אובדן דיוק.

FlashAttention-2 הוא המשך הגיוני לאלגוריתם FlashAttention, שהוצג על ידי אותה קבוצה בשנת 2022. הגרסה החדשה פותרת את בעיית העומס החלקי על ה-GPU שנצפתה אצל קודמו, ומשיגה עלייה של כמעט פי שניים במהירות בהשוואה לגרסה הראשונה.

רקע: בעיית ה-attention בטרנספורמרים

מנגנון ה-self-attention הסטנדרטי מהווה צוואר בקבוק בעבודה עם רצפי טקסט ארוכים בטרנספורמרים. המורכבות החישובית וצריכת הזיכרון שלו גדלות בריבועיות (O(N²)) בהתאם לאורך הרצף (N), מה שמטיל מגבלות חמורות על אורך ההקשר המרבי ועל יכולת ההרחבה של LLM[1].

לפתרון בעיה זו הוצג ב-2022 האלגוריתם FlashAttention[2]. רעיונותיו המרכזיים:

  • התחשבות בהיררכיית הזיכרון של GPU‏ (IO-awareness): האלגוריתם ממזער פעולות קריאה/כתיבה יקרות בין זיכרון ה-GPU האיטי (HBM) לבין זיכרון ה-SRAM המהיר על השבב.
  • עיבוד בבלוקים (tiling): החישובים מחולקים לבלוקים קטנים (tiles) המעובדים ב-SRAM המהיר, מה שמאפשר להימנע מהמטריאליזציה של מטריצת ה-attention המלאה בזיכרון.

הדבר אפשר להשיג צמיחה ליניארית בצריכת הזיכרון (O(N)) ותאוצה של פי 2–4 בהשוואה למימושים הסטנדרטיים[2]. FlashAttention זכה לאימוץ נרחב ותרם להופעתן של מודלים עם הקשר מוגדל משמעותית, למשל מ-2–4 אלף token‏ (GPT-3) עד 128 אלף (GPT-4) ומעלה[3]. כך, במודל Falcon-40B השימוש ב-FlashAttention האיץ את ה-inference פי 3, ואת ביצועי היצירה הכוללים — פי 5 בהשוואה ל-GPT-3[4].

פיתוח ומטרות FlashAttention-2

למרות ההצלחה, הגרסה הראשונה של FlashAttention לא ניצלה את משאבי החישוב של ה-GPU במלואם. בכרטיסי המסך NVIDIA A100 הביצועים הגיעו לכדי 25–40% בלבד מהמקסימום התיאורטי (FLOPs/s)[1]. הסיבה העיקרית הייתה עומס לא אופטימלי על מעבדי הזרמים (Streaming Multiprocessors) ופעולות עודפות עם הזיכרון המשותף[5].

מטרתו של FlashAttention-2 הייתה להאיץ עוד יותר את החישובים באמצעות מקבול יעיל יותר של העבודה ומזעור פעולות עזר. האלגוריתם נוסח מחדש במלואו תוך שימוש בפרימיטיבים ברמה נמוכה של ספריית NVIDIA CUTLASS 3.x להשגת ביצועים מרביים[6].

ארכיטקטורה טכנית ועקרונות עבודה

FlashAttention-2 מציג שלושה שיפורים מרכזיים לשיפור המקביליות והיעילות[1]:

1. מזעור פעולות שאינן מטריציות

האלגוריתם מצמצם את מספר פעולות הנקודה הצפה העזריות שאינן כפל מטריצות (non-matmul FLOPs). מאחר שגרעיני ה-tensor של ה-GPU מותאמים בדיוק לפעולות מטריציות (GEMM) ומבצעים אותן עד פי 16 מהר יותר, שינוי זה מאפשר לנצל את רוב הזמן בבלוקי ה-GPU היעילים ביותר.

2. מקביליות משופרת

ב-FlashAttention המקורי, עבודה על "ראש" attention אחד לא עברה מקבול, מה שגרם להשבתות בעת רצפים ארוכים וגדלי batch קטנים. FlashAttention-2 מציג מקביליות בין-בלוקית: כעת החישובים עבור ראש attention אחד מתחלקים בין מעבדי זרמים שונים של ה-GPU, מה שמגביר משמעותית את העומס עליהם.

3. חלוקת עבודה אופטימלית בתוך בלוק

ברמה של בלוק חישובי אחד, העבודה חולקה מחדש בין קבוצות threads‏ (warps) כדי לצמצם את החלפת הנתונים דרך הזיכרון המשותף (shared memory). הדבר מפחית את מספר פעולות הקריאה/כתיבה העודפות הנדרשות לנירמול Softmax.

ביצועים ויעילות

בזכות השיפורים הארכיטקטוניים, FlashAttention-2 מפגין עלייה משמעותית בביצועים:

  • תאוצה כפולה: האלגוריתם פועל בקירוב פי 2 מהר יותר בהשוואה לגרסה הראשונה של FlashAttention[1].
  • ניצול GPU גבוה: על GPU NVIDIA A100 מושגים 50–73% מרוחב הפס התיאורטי המרבי (TFLOPs), הקרוב ליעילות של פעולות כפל מטריצות אופטימליות (GEMM)[1].
  • מהירות חישוב שיא:
    • על GPU A100 מושגת מהירות של עד 225 TFLOP/s במחזור האימון המלא של מודל מסוג GPT, המתאים ל-72% ניצול של בלוקי החישוב. לשם השוואה, attention סטנדרטי באותם תנאים העמיס את ה-GPU בפחות מ-100 TFLOP/s[7].
    • על GPU H100 הביצועים מגיעים ל-335 TFLOP/s[7].

עלייה כזו בביצועים מאפשרת, למשל, לאמן מודל עם חלון הקשר של 16k token באותו זמן שנדרש בעבר לחלון 8k token[5]. חשוב לציין שהאלגוריתם נשאר מדויק ודטרמיניסטי, ולכן שימוש בו אינו משפיע על איכות תחזיות המודל[8].

שימוש ואינטגרציה במערכת האקולוגית

FlashAttention-2 הפך במהירות לכלי סטנדרטי במערכת האקולוגית של LLM. הוא משולב בפריימוורקים וספריות פופולריות רבות:

  • PyTorch: תמיכה מובנית.
  • Hugging Face Transformers: התמיכה מופעלת באמצעות הפרמטר `attn_implementation="flash_attention_2"` בעת טעינת המודל[9]. תואם לעשרות ארכיטקטורות (GPT,‏ Llama,‏ Falcon,‏ BERT ואחרות)[10].
  • TensorRT-LLM,‏ xFormers ו-Triton: האלגוריתם ממומש עבור פלטפורמות אלה, מה שמבטיח שימוש נרחב[7].

האינטגרציה מאפשרת לשלב בקלות את FlashAttention-2 עם שיטות אופטימיזציה אחרות, כגון קוונטיזציה (GPTQ,‏ QLoRA) ו-fine-tuning יעיל (PEFT)[9].

השוואה עם גרסאות עוקבות

FlashAttention-3

המחקר בתחום אופטימיזציית ה-attention נמשך. ביולי 2024 הציג טרי דאו את FlashAttention-3, המכוון לניצול יכולות ארכיטקטורת ה-GPU NVIDIA Hopper ‏(H100/H200). חידושים מרכזיים[3]:

  • תמיכה ב-FP8: שימוש בחישובי נקודה צפה של 8 סיביות להאצה נוספת.
  • פעולות אסינכרוניות: ניצול יעיל יותר של יכולות ה-GPU האסינכרוניות.

FlashAttention-3 מספק תאוצה של פי 1.5–2 בהשוואה ל-FlashAttention-2 על GPU H100, ומשיג ביצועים של עד 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_(HE)&oldid=2201"