FlashAttention-3: Attention nhanh hơn và chính xác hơn nhờ bất đồng bộ và độ chính xác thấp

• Nút thắt cổ chai Attention của Transformer từ lâu đã giới hạn tốc độ huấn luyện và suy luận của các LLM ngữ cảnh dài; FlashAttention-3 là phiên bản mới nhằm giảm mạnh hơn nữa nút thắt này bằng cách tận dụng các tính năng của GPU Hopper
• Trọng tâm là dùng tính bất đồng bộ của Tensor Core và TMA để chồng lấp tính toán với di chuyển dữ liệu, đồng thời bố trí xen kẽ matmul và softmax nhằm giảm thời gian GPU nhàn rỗi
• Trên H100, FlashAttention-2 chỉ khai thác được 35% FLOPS tối đa lý thuyết, còn FlashAttention-3 đạt tối đa 740 TFLOPS ở FP16, tương đương 75% mức lý thuyết của H100
• Ở độ chính xác thấp FP8, các giá trị ngoại lai trong activation của LLM làm tăng lỗi, nên FlashAttention-3 dùng incoherent processing dựa trên Hadamard transform để giảm lỗi 2,6 lần so với FP8 attention chuẩn
• FlashAttention-3 đã được công bố trên GitHub; khi chi phí Attention giảm, dư địa xử lý ngữ cảnh dài hơn và cải thiện hiệu quả huấn luyện/suy luận LLM sẽ lớn hơn

---

Mục tiêu và hiệu năng của FlashAttention-3

• Attention là lớp cốt lõi của Transformer, nhưng trở thành nút thắt cổ chai chính trong các mô hình ngôn ngữ lớn và ứng dụng ngữ cảnh dài
• FlashAttention và FlashAttention-2 đã tiên phong trong cách tiếp cận tăng tốc Attention bằng cách giảm đọc/ghi bộ nhớ GPU, và hiện được hầu hết thư viện sử dụng để tăng tốc huấn luyện và suy luận Transformer
• Cách tiếp cận này đã góp phần giúp độ dài ngữ cảnh LLM tăng trong 2 năm gần đây, từ 2–4K của GPT-3·OPT lên 128K của GPT-4 và 1M của Llama 3
• FlashAttention-2 chỉ đạt 35% mức sử dụng FLOPS tối đa lý thuyết trên GPU H100, trong khi FlashAttention-3 nâng mức này lên bằng cách tận dụng các tính năng mới của GPU Hopper
• Hiệu năng FlashAttention-3 được công bố như sau
  • Nhanh hơn FlashAttention-2 1,5–2,0 lần ở FP16
  • Tối đa 740 TFLOPS ở FP16
  • 75% mức sử dụng FLOPS tối đa lý thuyết của H100
  • Gần 1,2 PFLOPS ở FP8
  • Lỗi nhỏ hơn 2,6 lần so với FP8 attention chuẩn

Ôn lại cách FlashAttention hoạt động

• FlashAttention tăng tốc bằng cách sắp xếp lại thứ tự tính Attention, dùng tiling và tính toán lại, đồng thời giảm mức sử dụng bộ nhớ theo sequence length từ bậc hai xuống tuyến tính
• Nó tải các khối đầu vào từ HBM vào SRAM, thực hiện Attention trên khối đó, rồi cập nhật đầu ra vào HBM
• Vì không ghi ma trận Attention trung gian lớn vào HBM, lượng đọc/ghi bộ nhớ giảm, cho phép tăng tốc 2–4 lần theo thời gian chạy thực tế
• Khi kết hợp tiling với softmax rescaling, có thể xử lý theo từng khối mà vẫn thu được đầu ra đúng, không cần xấp xỉ

Tính năng GPU Hopper: WGMMA, TMA, FP8

• FlashAttention-2 có thể đạt tối đa 70% FLOPS lý thuyết trên GPU Ampere A100, nhưng chưa tận dụng đầy đủ các tính năng mới của GPU Hopper
• FlashAttention-3 sử dụng ba tính năng của Hopper
  • WGMMA: chức năng warpgroup matrix multiply-accumulate tận dụng Tensor Core mới của Hopper, có thông lượng cao hơn mma.sync của Ampere
  • TMA: đơn vị phần cứng chuyên dụng để tăng tốc truyền dữ liệu giữa global memory và shared memory, xử lý tính toán chỉ số và out-of-bound predication để giảm dùng register
  • FP8: có thể tăng gấp đôi thông lượng Tensor Core so với FP16, nhưng vì biểu diễn giá trị số thực dấu phẩy động bằng ít bit hơn nên có đánh đổi với độ chính xác
• FlashAttention-3 tận dụng các tính năng Hopper bằng cách dùng các abstraction của NVIDIA CUTLASS
• Chỉ riêng việc viết lại FlashAttention để dùng các tính năng mới đã cải thiện hiệu năng FP16 forward pass từ khoảng 350 TFLOPS của FlashAttention-2 lên mức 540–570 TFLOPS

Chồng lấp GEMM và softmax bằng bất đồng bộ

• Các phép toán chính của Attention là GEMM giữa Q-K, P-V và softmax
• Trên các bộ tăng tốc hiện đại, các phép toán không phải matmul chậm hơn matmul rất nhiều; các special function như exponential trong softmax được xử lý bởi đơn vị riêng, tách biệt với floating point multiply-add hoặc matrix multiply-add
• H100 SXM5 cung cấp 989 TFLOPS cho FP16 matrix multiply, nhưng thông lượng special function chỉ 3,9 TFLOPS, thấp hơn 256 lần
• Với head dimension 128, dù FLOPS của matmul nhiều hơn exponential 512 lần, exponential vẫn có thể chiếm 50% thời gian so với matmul
• Ở FP8, FLOPS của matmul nhanh gấp đôi nhưng tốc độ exponential không đổi, nên việc thực hiện song song matmul và softmax càng quan trọng hơn

• pingpong scheduling giữa các warpgroup

  • GPU warp scheduler tự động thực hiện một phần overlap bằng cách chạy các warp khác trong lúc một số warp chờ kết quả GEMM
  • FlashAttention-3 dùng synchronization barrier để chồng lấp GEMM và softmax của hai warpgroup tốt hơn theo cách thủ công
  • Warpgroup 1 trước tiên thực hiện GEMM1 của một iteration và GEMM0 của iteration tiếp theo
  • Sau đó, trong khi warpgroup 2 thực hiện GEMM, warpgroup 1 xử lý softmax
  • pingpong schedule này là cách ẩn softmax phía sau thời gian chạy GEMM của warpgroup khác
  • Lịch thực tế không hoàn toàn gọn gàng như sơ đồ, nhưng với FP16 attention forward pass ở head dimension 128 và sequence length 8K, nó nâng hiệu năng từ khoảng 570 TFLOPS lên 620 TFLOPS

• overlap bên trong warpgroup

  • Ngay cả trong một warpgroup, có thể chạy một phần softmax trong khi GEMM của warpgroup đó đang thực thi
  • Pipelining này nâng thông lượng FP16 attention forward từ khoảng 620 TFLOPS lên 640–660 TFLOPS
  • Đổi lại, register pressure tăng vì phải giữ đồng thời accumulator của GEMM và đầu vào/đầu ra của softmax
  • Nhìn tổng thể, kỹ thuật này đem lại tradeoff có lợi

Độ chính xác thấp FP8 và incoherent processing

• Activation của LLM có thể chứa các outlier có magnitude lớn hơn rất nhiều so với các feature còn lại
• Outlier khiến lượng tử hóa khó hơn và làm tăng mạnh quantization error
• FlashAttention-3 tận dụng incoherent processing được dùng trong tài liệu về lượng tử hóa như QuIP
• Bằng cách nhân query và key với một random orthogonal matrix, nó phân tán outlier và giảm lỗi lượng tử hóa
• Trong triển khai, FlashAttention-3 dùng Hadamard transform có random sign
  • Với head dimension là d, có thể thực hiện cho từng attention head trong thời gian O(d log d) thay vì O(d²)
  • Vì Hadamard transform bị giới hạn bởi memory bandwidth, có thể xử lý mà không tốn thêm chi phí nếu fuse với phép toán trước đó cũng bị giới hạn bởi memory bandwidth, như rotary embedding
• Trong thí nghiệm tạo Q, K, V từ phân phối chuẩn chuẩn hóa và đưa magnitude lớn vào 0,1% entry để mô phỏng outlier, incoherent processing giảm lỗi lượng tử hóa 2,6 lần

Benchmark và tình trạng công bố

• FlashAttention-3 được so sánh không chỉ với FlashAttention-2 mà còn với các triển khai Triton và cuDNN vốn đã dùng các tính năng phần cứng mới của GPU Hopper
• Ở FP16, nó cho thấy tốc độ tăng khoảng 1,6–1,8 lần so với FlashAttention-2
• Ở FP8, nó đạt gần 1,2 PFLOPS
• Kho GitHub FlashAttention-3 đã được công bố
• Bài báo cũng có thể xem trong cùng kho flash-attention

Các tối ưu còn lại và tích hợp trong tương lai

• Ngoài những nội dung được đề cập trong blog, bài báo còn bao gồm các tối ưu như variable length sequence, persistent kernel và in-kernel transpose cho FP8
• Thiết kế thuật toán phù hợp với phần cứng thực thi có thể đem lại cải thiện hiệu quả lớn và mở ra các khả năng mô hình mới như ngữ cảnh dài
• Công việc trong tương lai bao gồm tối ưu hóa suy luận LLM và khái quát hóa kỹ thuật sang các kiến trúc phần cứng khác
• FlashAttention-3 được kỳ vọng sẽ được tích hợp vào các bản phát hành PyTorch trong tương lai

[GN⁺]

Các ý kiến trên Hacker News

• Nhìn vào chú thích trong mã, có vẻ Tri Dao đã bắt tay vào FA3 từ tháng 4/2022, ngay sau khi Hopper/H100 được công bố
  Việc mất hơn 2 năm cho đến khi mã được công bố hôm nay hơi thú vị; có thể là vì những giải pháp tốt hơn đang được chuẩn bị
  Lịch sử bài báo gần đây của Tri nghiêng về SSM và các kiến trúc họ Mamba. FlashAttention có độ phức tạp thời gian bậc hai theo độ dài chuỗi, nhưng các thuật toán mới nhất là dưới bậc hai, nên không chỉ làm cùng phép tính hiệu quả hơn mà còn giảm đáng kể chính lượng tính toán
  Trong một bài báo dài năm nay, Dao và Gu đã cho thấy Mamba/SSM cũng có thể được hình thức hóa sao cho dễ tăng tốc bằng cùng các phép toán nguyên thủy phần cứng mà Transformer được hưởng lợi

  • Cho đến khi Giả thuyết thời gian mũ mạnh (SETH) được chứng minh hoặc bác bỏ, ta vẫn cần chi phí bậc hai hoặc phải từ bỏ một điều gì đó. Rốt cuộc đó là chi phí của tìm kiếm vét cạn
    Nếu chứng minh hoặc bác bỏ SETH thì cũng sẽ giải được bài toán P và NP, nên khó kỳ vọng chuyện đó sớm xảy ra
    Điểm cốt lõi là một trường hợp sử dụng cụ thể có chịu nổi chi phí đó hay không

• Tôi tò mò thuật toán FlashAttention bị ràng buộc với phần cứng đến mức nào
  Ví dụ, thông báo lần này nói rằng nó tận dụng tính năng bất đồng bộ của GPU H100, vậy có vẻ điều đó nghĩa là các card không thuộc dòng H sẽ không nhận được mức tăng tốc đó
  Ngoài ra, thư viện FlashAttention thực tế cần CUDA, nhưng thuật toán dường như đã được port sang Metal[^0]. Nếu thuật toán gần như là một hàm thuần thì chẳng phải nó nên có thể được triển khai trên bất kỳ GPU/framework machine learning nào sao
  [0]: https://github.com/philipturner/metal-flash-attention

  • Có nhiều câu trả lời hay, nhưng nói ngắn gọn thì trong thực tế nó bị ràng buộc với phần cứng “khá nhiều”. Dưới đây là một ví dụ ổn

    https://github.com/karpathy/nanoGPT/blob/master/model.py#L45
    nanoGPT của Karpathy kiểm tra sự tồn tại của torch.nn.functional.scaled_dot_product_attention để gọi FlashAttention
    https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html
    Nhìn vào tài liệu, trên thực tế bạn sẽ muốn nó gọi FA2 trong hầu hết trường hợp; FA2 tối ưu kernel của thiết bị để chia nhỏ phép Softmax trên ma trận tam giác và giảm việc đưa qua lại những batch số thực dấu phẩy động không cần thiết giữa GPU và CPU
    https://arxiv.org/pdf/2307.08691
    Bài báo FA2 gần như được trình bày hoàn toàn từ góc nhìn phần cứng mà nó chạy trên đó

  • Cải tiến thuật toán của FlashAttention chủ yếu là chia nhỏ và ghép lại phần Softmax của attention, và bản thân điều đó không phải là một ý tưởng hoàn toàn mới. Đóng góp áp đảo nằm ở việc triển khai phương pháp đó và các chi tiết liên quan một cách hiệu quả trên phần cứng Nvidia
  • Ban đầu FlashAttention hầu như không phụ thuộc vào phần cứng
    Các phiên bản mới thì tùy vào mức độ trừu tượng. ThunderKittens[0], tương tự những gì bài viết nói, đem lại mức tăng tốc khoảng 1,3–2 lần so với FA2, đồng thời vẫn áp dụng tương đối phổ quát trên nhiều GPU
    Với mỗi phần cứng mới có thể có các tính năng chuyên biệt cho phần cứng để khai thác thêm hiệu năng. Thường thì các nhà cung cấp sẽ áp dụng những tính năng giúp họ vượt lên trước, nhưng API và thư viện sẽ bị phân mảnh, như đã thấy với CUDA
    [0]: https://hazyresearch.stanford.edu/blog/2024-05-12-tk
  • Về mặt khái niệm thì có ràng buộc một chút, còn từ góc độ triển khai thực tế thì ràng buộc rất nhiều. Ngay cả triển khai Python tiêu chuẩn cũng bên trong biên dịch các kernel được điều chỉnh cho phần cứng cụ thể
  • Nói thêm từ góc nhìn thực tế, phần cứng AMD vẫn còn thiếu một triển khai flash-attention-2 đúng nghĩa. ROCm đang dần trở nên dùng được, nhưng vẫn chưa thể so sánh với CUDA

• Muốn hỏi các anh chị phía compiler. Liệu compiler có khả năng tự tìm ra các tối ưu hóa như FlashAttention không? TVM và tinygrad có vẻ đang đi theo hướng đó, nhưng khó tin là có thể làm được

  • Về mặt lý thuyết thì có thể. Nhờ các tính chất đại số của toán học, có thể tái sắp xếp ở quy mô lớn, rồi thêm vào đó tiling vòng lặp đa diện tương đối có cấu trúc là được
    Tuy nhiên chi phí lớn nên cần cache kết quả tìm kiếm đó
    Tối ưu hóa e-graph có vẻ rất hợp với lĩnh vực này. Nhưng nó đòi hỏi một chuyển đổi mô hình lớn trong cách xử lý các pass tối ưu hóa, nên ngoài một số công cụ ngách thì hầu như chưa được triển khai rộng rãi. Ví dụ, nó không khớp tốt với call graph truyền thống; để triển khai e-graph vượt ra ngoài/giữa các basic block và vòng lặp for thì phải thay đổi đáng kể luồng điều khiển, và break với return cũng không được hỗ trợ
  • Đây có vẻ là vấn đề cực kỳ khó nhưng không phải bất khả thi
    Tôi không rõ mức hiện đại nhất của tối ưu hóa compiler đã đi đến đâu về bố trí dữ liệu và tối đa hóa mức sử dụng bộ xử lý
    Trước đây tôi từng xem một video về tối ưu hóa nói rằng một tối ưu hóa nhỏ có giúp tăng tốc, nhưng ảnh hưởng đó không đáng kể so với biến động tốc độ đến từ khác biệt về bố trí bộ nhớ do chính tối ưu hóa đó, hoặc thậm chí các thay đổi ngẫu nhiên, tạo ra
    Bài trình bày đó tập trung hơn vào việc phân biệt tín hiệu trong nhiễu, nhưng bản thân nhiễu ấy là dấu vết cho thấy compiler còn chưa xử lý tốt cả những dạng đơn giản hơn nhiều so với vấn đề đang nói ở đây
    Chỉ riêng cấu trúc CPU–bộ nhớ, nơi cache và mẫu truy cập ảnh hưởng đến tốc độ, đã phức tạp rồi; nếu cộng thêm cả cấu trúc GPU thì có lẽ đây là vùng còn khá chưa được khai phá
    Một ngày nào đó có thể sẽ làm được. Vì đây là lĩnh vực AI, cũng nảy sinh câu hỏi liệu một AI đủ thông minh có làm được việc này không, nhưng còn tùy vào tiêu chuẩn của “đủ”
    Có thể hình dung một bài kiểm tra cấp rất cao cho mô hình AI: đưa cho nó thứ như micrograd và yêu cầu giữ nguyên cùng interface nhưng tạo ra thứ nhanh hơn torch. Hiện vẫn còn chưa gần đến mức đó, nhưng nếu làm được thì sẽ rất thú vị
  • Tôi nghĩ là không. Nên xem nó như một thuật toán khác. Thay vì chỉ xét toán học, đây là việc thiết kế thuật toán có tính đến hình dạng của phần cứng
    TVM thì có thể hiểu được. Nói nghiêm ngặt thì nó làm việc khác, nhưng là lĩnh vực khá gần
    Còn tinygrad thì tôi không rõ vì sao lại có cảm giác như vậy
  • https://github.com/uwplse/tensat
  • Việc gọi toán tử cấp cao từ một ngôn ngữ bọc như Python khá là khó

• Nếu có ai muốn port cái này sang ROCm / AMD MI300x thì hãy liên hệ [e-mail]. Tôi tuyệt đối không gửi spam
  Tôi có thể tài trợ thời gian tính toán cho việc này

  • Đây là công ty server accelerator AMD nhỉ! Công việc rất hay, hy vọng có ai đó nhận làm :)
  • Không có ý thô lỗ, nhưng tôi tò mò về mục đích của đề xuất này. Liệu có ai chỉ nhận quyền truy cập phần cứng rồi port việc này miễn phí không? Người đó được lợi gì?

• FlashAttention-3 is optimized for Hopper GPUs (e.g. H100).
  FA3 chạy trên GPU phổ thông như 3090 và 4090 thì hiệu năng ở mức nào?

  • Nó dành riêng cho Hopper. Các cải tiến gắn rất chặt với tính năng Hopper như warp group và TMA
    Trên 4090, có thể đạt tăng tốc nếu dùng bản triển khai Triton của FP8 attention: https://triton-lang.org/main/getting-started/tutorials/06-fused-attention.html

• TMA (Tensor Memory Accelerator). This is a special hardware unit that accelerates the transfer of data between global memory and shared memory, taking care of all index calculation and out-of-bound predication. This frees up registers, which is a valuable resource to increase tile size and efficiency.
  Theo tôi hiểu, TMA đúng là giúp giảm dùng register, nhưng điều quan trọng hơn là nó cho phần cứng xử lý tạo địa chỉ. Khi các phép toán xung quanh ngày càng nhanh, việc tạo địa chỉ có thể trở thành nút thắt

• Đây là một trong những cải tiến quan trọng nhất trong toàn bộ AI. Nó cho phép dùng được nhiều hơn và nhanh hơn trên cùng phần cứng, đồng thời đem lại lợi ích gần như không phải đánh đổi cho đa số người dùng AI

  • Với những người có H100 thì đúng vậy

• Tôi tò mò vì sao FlashAttention khi dùng masking biến đổi lại chậm hơn khoảng 5 lần so với khi không dùng. Nếu thiếu hỗ trợ masking tốt thì hiệu quả tối ưu hóa gần như biến mất

  • Bạn đang xem benchmark đó ở đâu vậy?

• Mong có chuyên gia trả lời giúp vài câu hỏi :)
  FlashAttention có phải là một thay thế drop-in cho phép tính attention của LLM không? Có thể dùng ở bất cứ đâu có dùng phép tính “attention” không, hay phải huấn luyện riêng LLM để dùng FA?
  FA có quan hệ thế nào với các chiến lược như GQA, sliding window attention? Chúng là các khái niệm trực giao với nhau, hay mỗi chiến lược cần một triển khai FA riêng?
  Gần đây llama.cpp đã thêm hỗ trợ FlashAttention; điều đó có nghĩa là nó bắt đầu dùng những thứ như CUDA kernel mà FlashAttention cung cấp không?
  Cuối cùng, bài này so sánh FlashAttention với Triton. Triton chẳng phải giống một lớp trừu tượng sao? Không thể triển khai FA bằng Triton à? Tôi chưa hiểu lắm cách nói “FlashAttention so với Triton”.

  • 1. Gần như đúng. Về mặt toán học là tương đương. Vấn đề ở cấp phần mềm chỉ là những thứ như quản lý phiên bản dependency hay định dạng dữ liệu trong bộ nhớ, và FlashAttention 2 đã có trong HuggingFace cùng nhiều thư viện phổ biến. FlashAttention 3 nhiều khả năng cũng sẽ sớm được đưa vào, nhưng để chạy thì cần GPU H100.
    2. FlashAttention 2 đã thêm hỗ trợ GQA trong một bản cập nhật trước đó:
       https://github.com/Dao-AILab/flash-attention
    3. Ở đây là so sánh triển khai FlashAttention này, được viết bằng CUDA C++ thuần, với triển khai Triton của một thuật toán tương tự được viết bằng Triton: https://triton-lang.org/main/getting-started/tutorials/06-fused-attention.html
  • FlashAttention có thể thay thế drop-in cho phép tính attention của LLM.
    FlashAttention là cách tính phần Softmax(QK^T)V của attention, còn GQA là cách tính các ma trận Q, K, V. Với sliding window attention thì tôi không chắc bằng, nhưng đó là cách thay đổi attention mask để kiểm soát query nào có thể nhìn thấy key nào.
    Tôi chưa dùng llama.cpp, nhưng cách giải thích rằng nó bắt đầu lấy CUDA kernel ra dùng nhìn chung có vẻ đúng.
    Câu hỏi cuối nói đến triển khai FlashAttention trước đây được viết bằng Triton.

• Tôi tò mò vì bài này nói các phép toán như sigmoid rất chậm.
  Các LLM hiện đại dùng rất nhiều hàm kích hoạt có sigmoid hoặc Softmax, như SiLU, Swish, SOLU.
  ReLU có ít tổn thất hiệu năng hơn không? Nếu vậy, liệu có nên quay lại ReLU tốt cũ không?

  • ReLU đúng nghĩa là một hàm tuyến tính bị cắt về 0 tại một điểm nào đó, nên lượng tính toán ít hơn rất nhiều so với những hàm có hàm mũ. Tuy nhiên, có lẽ khó đạt kết quả cạnh tranh với một hàm kích hoạt đơn giản như vậy.