FlashAttention (HE)

FlashAttention — הוא אלגוריתם פורץ דרך לחישוב מנגנון ה-attention, שפותח במטרה להאיץ באופן משמעותי את אימון ואינפרנס של מודלי שפה גדולים (LLM) תוך שמירה על דיוק חישובי מלא. האלגוריתם הוצג לראשונה בשנת 2022 על ידי צוות חוקרים מאוניברסיטת סטנפורד בהובלת טרי דאו (Tri Dao).

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

בעיות ה-attention הסטנדרטי

מנגנון ה-self-attention הסטנדרטי ב-transformer מחושב לפי הנוסחה: $\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{Q{K}^T}{\sqrt{d_k}}\right)V$ כאשר Q, K, V הן מטריצות השאילתות, המפתחות והערכים.

הבעיה העיקרית בגישה זו היא מורכבות ריבועית בזמן ובזיכרון (O(N²)) ביחס לאורך הרצף N. במימוש נאיבי יש לחשב ולשמור בזיכרון ה-GPU את מטריצת ה-attention המלאה S בגודל N×N, דבר המוביל לשתי בעיות קריטיות:

1. צריכת זיכרון גבוהה: שמירת מטריצה N×N הופכת לבלתי אפשרית בעבודה עם הקשרים ארוכים.
2. פעולות קלט-פלט (IO): צוואר הבקבוק העיקרי אינו מספר הפעולות האריתמטיות, אלא הגישות החוזרות ונשנות לזיכרון האיטי של ה-GPU.

היררכיית זיכרון GPU

להבנת הבעיה חשוב להבחין בין שני סוגי זיכרון ב-GPU (בדוגמת NVIDIA A100):

• SRAM (זיכרון סטטי): זיכרון מהיר על-שבב בנפח קטן (~20 MB) עם רוחב פס עצום (עד 19 TB/s).
• HBM (זיכרון בעל רוחב פס גבוה): זיכרון איטי בנפח גדול (40–80 GB) עם רוחב פס נמוך בהרבה (כ-1.5 TB/s).

אסימטריה זו הופכת את אלגוריתם ה-attention הסטנדרטי למוגבל רוחב פס זיכרון (memory-bound), שכן הוא קורא וכותב ללא הרף מטריצות גדולות מה-HBM האיטי, וזהו המקור העיקרי לעיכובים.

חידושי המפתח של FlashAttention

FlashAttention הוא אלגוריתם מודע-IO (IO-aware), הפותר את הבעיה על ידי מזעור הגישות ל-HBM. הדבר מושג באמצעות שלוש טכניקות עיקריות.

Tiling ועיבוד בלוקי

במקום עיבוד המטריצה כולה בבת אחת, FlashAttention מפצל את מטריצות הקלט Q, K, V לבלוקים קטנים (tiles) שמתאימים ל-SRAM המהיר. האלגוריתם טוען בלוקים אלה ברצף, מבצע עבורם את כל חישובי ה-attention ומעדכן את התוצאה הסופית, מבלי לשמור את מטריצת ה-attention המלאה ב-HBM האיטי.

חישוב Softmax מקוון

פריצת הדרך הטכנית המרכזית הייתה חישוב Softmax "מקוון" (online). ה-Softmax הסטנדרטי דורש הכרת כל אלמנטי וקטור הקלט לצורך נרמול. FlashAttention משתמש באלגוריתם מותאם המאפשר חישוב Softmax בחלקים. הוא שומר שני ערכים ביניים (המקסימום הנוכחי וסכום האקספוננטים), המתעדכנים עם עיבוד בלוקים חדשים, מה שמאפשר לקבל תוצאה מדויקת ללא גישה למטריצה כולה בבת אחת.

מיזוג פעולות לליבת CUDA אחת

כל פעולות ה-attention (כפל המטריצות QKᵀ, מיסוך, Softmax, כפל ב-V) משולבות בליבת CUDA ממוזגת אחת (fused kernel). פעולה זו מצמצמת דרמטית את מספר פעולות הקריאה/כתיבה ב-HBM: במקום מעברים חוזרים על פני המטריצה כולה, האלגוריתם טוען בלוק ל-SRAM פעם אחת, מבצע את כל החישובים וכותב רק את התוצאה הסופית.

יעילות תיאורטית ומעשית

מורכבות ואופטימליות

FlashAttention מצמצם את צריכת הזיכרון מ-O(N²) ל-O(N), המבטיח קנה מידה ליניארי. הוכח כי מורכבות ה-IO של האלגוריתם היא אופטימלית תיאורטית לחישוב attention בהיררכיית זיכרון דו-שכבתית, כלומר לא ניתן לבצע attention מדויק במהירות גבוהה יותר ללא שינוי חומרה.

תוצאות אמפיריות

הגרסה הראשונה של FlashAttention הדגימה שיפורים משמעותיים:

• האצה:
  • BERT-large (אורך 512): האצת אימון של 15%.
  • GPT-2 (אורך 1K): האצה פי 3.
  • משימות Long-Range Arena (1K–4K): האצה פי 2.4.
• חיסכון בזיכרון: עד חיסכון פי 20 בזיכרון בהשוואה למימושי הבסיס המדויקים.
• שיפור באיכות המודלים: הודות ליכולת לעבוד עם הקשרים ארוכים יותר, FlashAttention לא רק שאינו מאבד איכות, אלא אף משפר אותה. לדוגמה, הפרפלקסיה של GPT-2 השתפרה ב-0.7 נקודות, והדיוק במשימות סיווג מסמכים ארוכים עלה ב-6.4 נקודות.

התפתחות ופיתוחים נוספים

הצלחת FlashAttention הניחה את היסוד לסדרה שלמה של אלגוריתמים מוכווני-חומרה.

FlashAttention-2 (2023)

הגרסה השנייה נועדה לניצול מלא יותר של משאבי ה-GPU. ב-FlashAttention המקורי, היעילות על NVIDIA A100 עמדה על 25–40% בלבד מהמקסימום. FlashAttention-2 הציגה שיפורים בתיקבול החישובים, מה שאפשר:

• השגת האצה פי שניים בהשוואה לגרסה הראשונה.
• הגדלת ניצולת ה-GPU ל-50–73% מהמקסימום התיאורטי.
• הרחבת התמיכה לראשי attention בגודל 256, וכן לארכיטקטורות Multi-Query Attention (MQA).

FlashAttention-3 (2024)

הגרסה השלישית אוּפטמה במיוחד לארכיטקטורת GPU NVIDIA Hopper (H100). היא מנצלת יכולות חומרה חדשות, כגון אסינכרוניות Tensor Cores ותמיכה ב-FP8, מה שאפשר:

• השגת האצה נוספת פי 1.5–2 בהשוואה ל-FlashAttention-2.
• השגת ביצועים של עד 740 TFLOPS על FP16 וקרוב ל-1.2 PFLOPS על FP8.

פתרונות מיוחדים

רעיונות FlashAttention פותחו בפרויקטים נוספים:

• FlashInfer (2025): מנוע attention מותאם-אישית, אוּפטם במיוחד למשימות אינפרנס של LLM. הוא מתמקד בעבודה יעילה עם KV-cache במצב יצירה זורמת.
• FlashMLA (2024): מימוש attention עם דחיסת מטמון ההקשר (latent attention), המאפשר חיסכון בזיכרון על רצפים ארוכים מאוד עם אובדן מינימלי של מידע.

השפעה על התעשייה ועל המערכת האקולוגית

FlashAttention הפך לפריצת דרך יסודית והתגבש במהירות לסטנדרט תעשייתי לאימון ואינפרנס יעיל של LLM. הוא שולב בספריות מרכזיות כגון PyTorch ו-Hugging Face, ומשמש את רוב מודלי השפה הגדולים (LLaMA, MPT, Falcon, Claude ועוד).

דווקא FlashAttention וגרסאותיו המאוחרות מילאו תפקיד מכריע בהגדלת חלונות ההקשר של מודלי השפה: מ-2–4 אלף token (GPT-3) עד 128 אלף token (GPT-4) ואף עד מיליוני token במודלים ניסיוניים. האלגוריתם ביטל אחד המחסומים העיקריים על דרך הרחבת ה-transformer, ופתח אפשרויות חדשות ליישומי בינה מלאכותית — מניתוח מסמכים ארוכים ועד להבנה מולטי-מודלית.

קישורים

• המאגר הרשמי של FlashAttention ב-GitHub

ביבליוגרפיה

• Dao, T. et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. arXiv:2205.14135.
• Dao, T. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. arXiv:2307.08691.
• 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.
• Hong, K. et al. (2023). FlashDecoding++: Faster Large Language Model Inference on GPUs. arXiv:2311.01282.
• Ye, Z. et al. (2025). FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving. arXiv:2501.01005.
• 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. (2025). FlashMask: Efficient and Rich Mask Extension of FlashAttention. OpenReview wUtXB43Chi.
• Dao, T. et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (OpenReview version). OpenReview H4DqfPSibmx.
• Gholami, A. et al. (2024). FlashAttention on a Napkin: A Diagrammatic Approach to Deep Learning IO-Awareness. OpenReview pF2ukh7HxA.