FlashAttention-3: Attention ที่เร็วและแม่นยำยิ่งขึ้นด้วยความไม่ซิงโครนัสและความละเอียดต่ำ

 (together.ai)
1 คะแนน โดย GN⁺ 2024-07-12 | 1 ความคิดเห็น | แชร์ทาง WhatsApp
  • คอขวดของ Attention ใน Transformer ได้จำกัดความเร็วการฝึกและการอนุมานของ LLM ที่มีคอนเท็กซ์ยาวมาโดยตลอด และ FlashAttention-3 คือเวอร์ชันใหม่ที่มุ่งลดปัญหานี้ลงให้มากขึ้นโดยใช้ความสามารถของ GPU Hopper
  • แกนสำคัญคือการใช้ ความไม่ซิงโครนัส ของ Tensor Core และ TMA เพื่อซ้อนทับการคำนวณกับการย้ายข้อมูล และจัดวาง matmul กับ softmax แบบสลับกันเพื่อลดเวลาว่างของ GPU
  • บน H100, FlashAttention-2 ใช้ประโยชน์จาก FLOPS เชิงทฤษฎีสูงสุดได้เพียง 35% แต่ FlashAttention-3 ไปได้ถึง 740 TFLOPS ใน FP16 หรือ 75% ของค่าทฤษฎีของ H100
  • ในโหมดความละเอียดต่ำ FP8 เนื่องจาก outlier ใน activation ของ LLM ทำให้ความผิดพลาดเพิ่มขึ้น จึงใช้ incoherent processing ที่อิงกับ Hadamard transform เพื่อลดความผิดพลาดลง 2.6 เท่าเมื่อเทียบกับ FP8 attention แบบอ้างอิง
  • FlashAttention-3 เปิดเผยบน GitHub แล้ว และยิ่งต้นทุนของ Attention ลดลงมากเท่าไร ก็ยิ่งมีโอกาสเพิ่มประสิทธิภาพการรองรับคอนเท็กซ์ที่ยาวขึ้น รวมถึงการฝึกและการอนุมานของ LLM ได้มากขึ้น

เป้าหมายและประสิทธิภาพของ FlashAttention-3

  • Attention เป็นเลเยอร์หลักของ Transformer แต่กลายเป็น คอขวดสำคัญ ในโมเดลภาษาขนาดใหญ่และแอปพลิเคชันที่ใช้คอนเท็กซ์ยาว
  • FlashAttention และ FlashAttention-2 เป็นผู้บุกเบิกแนวทางเร่ง Attention ด้วยการลดการอ่าน·เขียนหน่วยความจำของ GPU และปัจจุบันไลบรารีส่วนใหญ่ก็นำไปใช้เพื่อเร่งการฝึกและการอนุมานของ Transformer
  • แนวทางนี้มีส่วนช่วยให้ความยาวคอนเท็กซ์ของ LLM เพิ่มขึ้นจาก 2~4K ของ GPT-3·OPT ในช่วง 2 ปีที่ผ่านมา ไปเป็น 128K ของ GPT-4 และ 1M ของ Llama 3
  • FlashAttention-2 ใช้ประโยชน์จาก FLOPS เชิงทฤษฎีสูงสุดของ GPU H100 ได้เพียง 35% แต่ FlashAttention-3 ดึงความสามารถนี้ขึ้นไปด้วยฟีเจอร์ใหม่ของ GPU Hopper
  • ประสิทธิภาพที่เปิดเผยของ FlashAttention-3 มีดังนี้
    • ใน FP16 เร็วกว่า FlashAttention-2 1.5~2.0 เท่า
    • ใน FP16 ได้สูงสุด 740 TFLOPS
    • ใช้ประโยชน์จาก FLOPS เชิงทฤษฎีสูงสุดของ H100 ได้ 75%
    • ใน FP8 ได้เกือบ 1.2 PFLOPS
    • มี ความผิดพลาดน้อยกว่า 2.6 เท่า เมื่อเทียบกับ FP8 attention แบบอ้างอิง

ทบทวนแนวทางของ FlashAttention

  • FlashAttention เร่งความเร็วด้วยการจัดลำดับการคำนวณ Attention ใหม่ และใช้ tiling กับการคำนวณซ้ำ พร้อมลดการใช้หน่วยความจำตาม sequence length จาก quadratic เหลือ linear
  • ระบบจะโหลดบล็อกอินพุตจาก HBM ไปยัง SRAM ทำ Attention กับบล็อกนั้น แล้วอัปเดตเอาต์พุตกลับไปยัง HBM
  • เนื่องจากไม่ต้องเขียนเมทริกซ์ Attention ขนาดใหญ่ระหว่างทางลง HBM จึงลดการอ่าน·เขียนหน่วยความจำได้ และทำให้เวลาในการรันจริงเร็วขึ้นได้ 2~4 เท่า
  • เมื่อใช้ tiling ร่วมกับ softmax rescaling ก็สามารถประมวลผลแบบเป็นบล็อกและยังได้เอาต์พุตที่ถูกต้องโดยไม่ต้องประมาณค่า

ความสามารถของ GPU Hopper: WGMMA, TMA, FP8

  • FlashAttention-2 สามารถไปได้ถึง 70% ของ FLOPS เชิงทฤษฎีสูงสุดบน GPU Ampere A100 แต่ยังใช้ฟีเจอร์ใหม่ของ GPU Hopper ได้ไม่เต็มที่
  • FlashAttention-3 ใช้ 3 ฟีเจอร์ของ Hopper
    • WGMMA: ความสามารถ warpgroup matrix multiply-accumulate ที่ใช้ Tensor Core รุ่นใหม่ของ Hopper และมี throughput สูงกว่า mma.sync ของ Ampere
    • TMA: หน่วยฮาร์ดแวร์เฉพาะสำหรับเร่งการส่งข้อมูลระหว่าง global memory กับ shared memory โดยจัดการการคำนวณ index และ out-of-bound predication เพื่อลดการใช้ register
    • FP8: เพิ่ม throughput ของ Tensor Core ได้เป็นสองเท่าเมื่อเทียบกับ FP16 แต่ก็มี tradeoff ด้านความแม่นยำเพราะใช้บิตน้อยลงในการแทนค่าทศนิยมลอยตัว
  • FlashAttention-3 ใช้ abstraction ของ NVIDIA CUTLASS เพื่อดึงความสามารถของ Hopper มาใช้
  • เพียงแค่เขียน FlashAttention ใหม่ให้ใช้ฟีเจอร์เหล่านี้ ประสิทธิภาพของ FP16 forward pass ก็เพิ่มจากราว 350 TFLOPS ของ FlashAttention-2 ไปเป็น 540~570 TFLOPS

ซ้อนทับ GEMM กับ softmax ด้วยความไม่ซิงโครนัส

  • การคำนวณหลักของ Attention คือ GEMM ระหว่าง Q-K และ P-V รวมถึง softmax
  • บนตัวเร่งความเร็วสมัยใหม่ การคำนวณที่ไม่ใช่ matmul ช้ากว่า matmul มาก และ special function อย่าง exponential ใน softmax ก็ถูกประมวลผลด้วยยูนิตแยกจาก floating point multiply-add หรือ matrix multiply-add
  • H100 SXM5 ให้ 989 TFLOPS สำหรับ FP16 matrix multiply แต่ throughput ของ special function อยู่ที่ 3.9 TFLOPS ต่ำกว่าถึง 256 เท่า
  • เมื่อ head dimension เท่ากับ 128 ถึงแม้ matmul FLOPS จะมากกว่า exponential 512 เท่า แต่ exponential ก็ยังอาจกินเวลาถึง 50% ของ matmul ได้
  • ใน FP8 แม้ matmul FLOPS จะเร็วขึ้นเป็นสองเท่า แต่ความเร็วของ exponential ยังเท่าเดิม ทำให้การรัน matmul กับ softmax แบบขนานกันยิ่งสำคัญขึ้น

  • pingpong scheduling ระหว่าง warpgroup

    • ตัวจัดตาราง warp ของ GPU สามารถรัน warp อื่นได้โดยอัตโนมัติระหว่างที่บาง warp กำลังรอผล GEMM ทำให้เกิด overlap บางส่วน
    • FlashAttention-3 ใช้ synchronization barrier เพื่อซ้อนทับ GEMM กับ softmax ของสอง warpgroup ได้ดียิ่งขึ้นแบบกำหนดเอง
    • warpgroup 1 จะรัน GEMM1 ของ iteration หนึ่งและ GEMM0 ของ iteration ถัดไปก่อน
    • จากนั้นระหว่างที่ warpgroup 2 รัน GEMM, warpgroup 1 จะประมวลผล softmax
    • ตารางแบบ pingpong นี้ซ่อนเวลา softmax ไว้หลังเวลารัน GEMM ของ warpgroup อื่น
    • แม้การจัดตารางจริงจะไม่เรียบร้อยสมบูรณ์แบบเหมือนในภาพ แต่ก็ทำให้ FP16 attention forward pass สำหรับ head dimension 128 และ sequence length 8K เพิ่มจากราว 570 TFLOPS เป็น 620 TFLOPS

  • overlap ภายใน warpgroup

    • ภายใน warpgroup เดียวกัน ก็สามารถรัน softmax บางส่วนระหว่างที่ GEMM ของ warpgroup นั้นกำลังทำงานได้
    • การทำ pipelining นี้เพิ่ม throughput ของ FP16 attention forward จากราว 620 TFLOPS ไปเป็น 640~660 TFLOPS
    • แต่ก็ทำให้ register pressure เพิ่มขึ้น เพราะต้องเก็บทั้ง GEMM accumulator และอินพุต·เอาต์พุตของ softmax ไว้พร้อมกัน
    • โดยรวมแล้วเทคนิคนี้ให้ tradeoff ที่คุ้มค่า

FP8 ความละเอียดต่ำและ incoherent processing

  • activation ของ LLM อาจมี outlier ที่มี magnitude สูงกว่าฟีเจอร์อื่นมาก
  • outlier ทำให้การ quantization ยากขึ้นและเพิ่ม quantization error อย่างมาก
  • FlashAttention-3 ใช้ incoherent processing ที่ถูกใช้ในงานวิจัยด้าน quantization เช่น QuIP
  • ระบบจะคูณ query และ key ด้วย random orthogonal matrix เพื่อกระจาย outlier และลดความผิดพลาดจากการ quantization
  • ในการนำไปใช้จริง จะใช้ Hadamard transform ที่มี random sign
    • หาก head dimension คือ d จะสามารถทำงานกับแต่ละ attention head ได้ในเวลา O(d log d) ไม่ใช่ O(d²)
    • Hadamard transform มีลักษณะเป็น memory-bandwidth bound ดังนั้นหากรวมเข้ากับการคำนวณก่อนหน้าอย่าง rotary embedding ที่ก็เป็น memory-bandwidth bound เช่นกัน ก็จะทำได้โดยแทบไม่มีต้นทุนเพิ่ม
  • ในการทดลองที่สร้าง Q, K, V จากการแจกแจงปกติมาตรฐาน และใส่ค่า magnitude สูงลงใน 0.1% ของ entry เพื่อจำลอง outlier, incoherent processing ลดความผิดพลาดจากการ quantization ลงได้ 2.6 เท่า

เบนช์มาร์กและสถานะการเปิดเผย

  • FlashAttention-3 ไม่ได้เปรียบเทียบแค่กับ FlashAttention-2 แต่ยังเทียบกับการติดตั้งบน Triton และ cuDNN ที่ใช้ฟีเจอร์ฮาร์ดแวร์ใหม่ของ Hopper อยู่แล้วด้วย
  • ใน FP16 มีความเร็วเพิ่มขึ้นประมาณ 1.6~1.8 เท่า เมื่อเทียบกับ FlashAttention-2
  • ใน FP8 ไปได้เกือบ 1.2 PFLOPS
  • คลัง GitHub ของ FlashAttention-3 เปิดเผยแล้ว
  • สามารถดูบทความวิจัยได้จาก คลัง flash-attention เดียวกัน

การปรับแต่งที่ยังเหลือและการรวมเข้าระบบในอนาคต

  • ในบทความวิจัยยังมีการปรับแต่งอื่นนอกเหนือจากที่กล่าวในบล็อก เช่น variable length sequence, persistent kernel และ in-kernel transpose สำหรับ FP8
  • การออกแบบอัลกอริทึมให้เหมาะกับฮาร์ดแวร์ที่ใช้สามารถสร้างการเพิ่มประสิทธิภาพอย่างมาก และเปิดทางให้กับความสามารถใหม่ของโมเดล เช่น คอนเท็กซ์ที่ยาวขึ้น
  • งานในอนาคตรวมถึงการปรับแต่ง LLM inference และการทำให้เทคนิคนี้ใช้ได้กับสถาปัตยกรรมฮาร์ดแวร์อื่น
  • คาดว่า FlashAttention-3 จะถูกรวมเข้าในรีลีส PyTorch ในอนาคต

บทความที่เกี่ยวข้อง

1 ความคิดเห็น

 
GN⁺ 2024-07-12
ความคิดเห็นจาก Hacker News

  • ดูจากคอมเมนต์ในโค้ดแล้ว ดูเหมือนว่า Tri Dao เริ่มทำงานกับ FA3 มาตั้งแต่ เดือนเมษายน 2022 ซึ่งเป็นช่วงไม่นานหลังการเปิดตัว Hopper/H100
    น่าสนใจเล็กน้อยที่โค้ดใช้เวลากว่า 2 ปีจึงถูกเปิดเผยในวันนี้ อาจเป็นเพราะกำลังมีวิธีแก้ปัญหาที่ดีกว่าเตรียมอยู่ก็ได้
    ประวัติผลงานวิจัยล่าสุดของ Tri เอนเอียงไปทางโครงสร้างตระกูล SSM และ Mamba มากขึ้น FlashAttention มี ความซับซ้อนเวลาเชิงกำลังสอง ตามความยาวลำดับ แต่อัลกอริทึมล่าสุดเป็นแบบต่ำกว่ากำลังสอง จึงไม่ใช่แค่ทำงานคำนวณเดิมได้มีประสิทธิภาพขึ้น แต่ลดปริมาณการคำนวณลงได้มากตั้งแต่ต้น
    ปีนี้ Dao และ Gu แสดงไว้ในบทความวิจัยขนาดยาวว่า Mamba/SSM ก็สามารถจัดรูปแบบอย่างเป็นทางการให้เหมาะกับการเร่งความเร็วด้วย primitive ของฮาร์ดแวร์แบบเดียวกับที่ Transformer ได้ประโยชน์ได้เช่นกัน

    • จนกว่า Strong Exponential Time Hypothesis (SETH) จะถูกพิสูจน์หรือหักล้าง ก็ยังต้องใช้ต้นทุนเชิงกำลังสอง หรือไม่ก็ต้องยอมสละบางอย่าง สุดท้ายแล้วนี่คือต้นทุนของการค้นหาแบบ exhaustive search
      หากพิสูจน์หรือหักล้าง SETH ได้ ก็จะแก้ปัญหา P กับ NP ได้ด้วย ดังนั้นจึงคาดหวังให้เกิดขึ้นเร็ว ๆ นี้ได้ยาก
      ประเด็นสำคัญคือกรณีใช้งานเฉพาะนั้นรับต้นทุนนี้ไหวหรือไม่

  • สงสัยว่าอัลกอริทึม FlashAttention ผูกกับฮาร์ดแวร์มากแค่ไหน
    เช่น ในประกาศครั้งนี้บอกว่าใช้ประโยชน์จาก ความสามารถแบบอะซิงโครนัส ของ GPU H100 ซึ่งดูเหมือนหมายความว่าการ์ดที่ไม่ใช่ซีรีส์ H จะไม่ได้รับความเร็วที่เพิ่มขึ้นนี้
    อีกอย่าง ไลบรารี FlashAttention ตัวจริงต้องใช้ CUDA แต่ดูเหมือนว่าอัลกอริทึมจะถูกพอร์ตไป Metal แล้ว[^0] หากอัลกอริทึมใกล้เคียงกับฟังก์ชันบริสุทธิ์ ก็น่าจะนำไปใช้งานได้บน GPU/เฟรมเวิร์กแมชชีนเลิร์นนิงใด ๆ ไม่ใช่หรือ
    [0]: https://github.com/philipturner/metal-flash-attention

    • มีคำตอบดี ๆ หลายข้อ แต่ถ้าพูดสั้น ๆ คือ “ในทางปฏิบัติ ค่อนข้างมาก” ที่ผูกกับฮาร์ดแวร์ ตัวอย่างด้านล่างค่อนข้างดี

      https://github.com/karpathy/nanoGPT/blob/master/model.py#L45
      nanoGPT ของ Karpathy ตรวจสอบว่ามี torch.nn.functional.scaled_dot_product_attention อยู่หรือไม่ แล้วจึงเรียก FlashAttention
      https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html
      หากดูเอกสาร จะเห็นว่าในทางปฏิบัติส่วนใหญ่ต้องการให้เรียก FA2 และ FA2 จะปรับแต่งเคอร์เนลของอุปกรณ์เพื่อแบ่งการทำงาน Softmax ของเมทริกซ์สามเหลี่ยม และลดการย้ายชุดข้อมูลทศนิยมลอยตัวที่ไม่จำเป็นไปมาระหว่าง GPU กับ CPU
      https://arxiv.org/pdf/2307.08691
      บทความวิจัย FA2 แทบทั้งหมดเขียนอธิบายจากมุมมองของฮาร์ดแวร์ที่มันรันอยู่

    • การปรับปรุงเชิงอัลกอริทึมของ FlashAttention หลัก ๆ คือ การแบ่งและรวมส่วน Softmax ของ attention ซึ่งตัวแนวคิดเองไม่ใช่เรื่องใหม่ทั้งหมด คุณูปการที่โดดเด่นคือการนำวิธีนั้นและรายละเอียดต่าง ๆ ไปใช้งานบนฮาร์ดแวร์ Nvidia ได้อย่างมีประสิทธิภาพ
    • เดิมที FlashAttention แทบไม่มีการพึ่งพาฮาร์ดแวร์
      เวอร์ชันล่าสุดขึ้นอยู่กับระดับการนามธรรม ThunderKittens[0] ให้ความเร็วเพิ่มขึ้นประมาณ 1.3~2 เท่าเมื่อเทียบกับ FA2 ตามที่บทความกล่าวไว้ และยังค่อนข้างใช้งานได้ทั่วไปกับ GPU หลากหลายตัว
      ฮาร์ดแวร์ใหม่แต่ละรุ่นอาจมีฟีเจอร์เฉพาะฮาร์ดแวร์ที่ดึงประสิทธิภาพเพิ่มได้ โดยทั่วไปผู้ผลิตมักรับฟีเจอร์ที่ช่วยให้ตนเหนือกว่าเข้ามา แต่ API และไลบรารีก็จะแตกกระจายเหมือนที่เกิดขึ้นกับ CUDA แล้ว
      [0]: https://hazyresearch.stanford.edu/blog/2024-05-12-tk
    • ในเชิงแนวคิดผูกอยู่บ้าง แต่ในมุมการใช้งานจริงผูกอยู่มาก แม้แต่ implementation มาตรฐานใน Python ก็ยังคอมไพล์ เคอร์เนล ที่ปรับให้เข้ากับฮาร์ดแวร์เฉพาะภายใน
    • ถ้าเสริมจากมุมปฏิบัติจริง ฮาร์ดแวร์ AMD ยังขาด implementation ที่ดีของ flash-attention-2 อยู่ ROCm ค่อย ๆ ใช้งานได้ดีขึ้น แต่ยังไม่ถึงระดับที่จะเทียบกับ CUDA ได้

  • อยากถามคนฝั่งคอมไพเลอร์ว่า มีความเป็นไปได้ไหมที่คอมไพเลอร์จะค้นพบการปรับแต่งประสิทธิภาพแบบ FlashAttention ได้เอง TVM กับ tinygrad ดูเหมือนกำลังไปในทิศทางนั้น แต่ก็ยากที่จะเชื่อว่าเป็นไปได้

    • ในทางทฤษฎีเป็นไปได้ คุณสมบัติทางพีชคณิตของคณิตศาสตร์ทำให้สามารถจัดเรียงใหม่ได้ในระดับใหญ่ และเติม polyhedral loop tiling ที่ค่อนข้างเป็นรูปแบบตายตัวเข้าไปก็พอ
      แต่ต้นทุนสูง จึงต้องแคชผลลัพธ์จากการค้นหานั้นไว้
      การปรับแต่งด้วย e-graph ดูเหมาะกับพื้นที่นี้ดี เพียงแต่ต้องเปลี่ยนกระบวนทัศน์ครั้งใหญ่ในวิธีจัดการ optimization pass จึงแทบไม่ได้ถูกนำไปใช้งานจริงนอกจากเครื่องมือเฉพาะทางบางส่วน เช่น มันไม่เข้ากับ call graph แบบดั้งเดิมนัก หากจะ deploy e-graph ให้ข้ามออกไปนอก/ระหว่าง basic block และ for loop ก็ต้องเปลี่ยน control flow อย่างมาก และยังไม่รองรับ break กับ return ด้วย
    • เป็นปัญหาที่ยากสุด ๆ แต่ดูไม่ใช่ว่าเป็นไปไม่ได้
      ไม่แน่ใจว่าระดับล่าสุดของการปรับแต่งคอมไพเลอร์ไปถึงไหนแล้วในแง่ของ การจัดวางข้อมูล และการใช้ประโยชน์จากโปรเซสเซอร์ให้สูงสุด
      เคยเห็นในวิดีโอเกี่ยวกับการปรับแต่งสมัยก่อนว่า optimization เล็ก ๆ ช่วยเพิ่มความเร็วได้ก็จริง แต่ผลกระทบนั้นน้อยมากเมื่อเทียบกับความผันผวนของความเร็วที่มาจากความต่างของการจัดวางหน่วยความจำ ซึ่งเกิดจาก optimization นั้นหรือแม้แต่การเปลี่ยนแบบสุ่ม
      งานนำเสนอนั้นเน้นไปที่การแยกสัญญาณออกจาก noise มากกว่า แต่ noise นั้นเองเป็นร่องรอยว่าคอมไพเลอร์ยังจัดการแม้แต่รูปแบบที่ง่ายกว่าปัญหาที่พูดถึงตรงนี้มากได้ไม่ดีนัก
      แค่สถาปัตยกรรม CPU·หน่วยความจำที่ cache และ access pattern ส่งผลต่อความเร็วก็ซับซ้อนแล้ว พอเพิ่มโครงสร้าง GPU เข้าไปอีก ก็น่าจะเป็นพื้นที่ที่ยังบุกเบิกไปไม่มาก
      สักวันหนึ่งอาจเป็นไปได้ เพราะเป็นวงการ AI ก็เลยเกิดคำถามว่า AI ที่ฉลาดพอจะทำสิ่งนี้ได้ไหม แต่ก็ขึ้นอยู่กับเกณฑ์ของคำว่า “พอ”
      ในฐานะการทดสอบระดับสูงมากของโมเดล AI อาจนึกถึงการให้สิ่งอย่าง micrograd แล้วสั่งให้สร้างสิ่งที่เร็วกว่า torch โดยคงอินเทอร์เฟซเดิมไว้ ตอนนี้ยังไปไม่ถึงแถว ๆ นั้นเลย แต่ถ้าเป็นไปได้ก็คงน่าสนใจ
    • คิดว่าไม่ใช่ ควรมองเหมือนเป็น อัลกอริทึมคนละแบบ คือออกแบบอัลกอริทึมโดยคำนึงถึงรูปร่างของฮาร์ดแวร์ แทนที่จะดูแค่คณิตศาสตร์
      TVM นั้นเข้าใจได้ พูดให้เคร่งครัดคือมันทำคนละอย่าง แต่เป็นพื้นที่ที่ค่อนข้างใกล้กัน
      แต่ tinygrad นี่ไม่รู้ว่าทำไมถึงรู้สึกแบบนั้น
    • https://github.com/uwplse/tensat
    • การเรียก โอเปอเรเตอร์ระดับสูง จากภาษาที่ครอบไว้อย่าง Python นั้นค่อนข้างยุ่งยาก

  • ถ้ามีใครอยากพอร์ตสิ่งนี้ไปยัง ROCm / AMD MI300x ติดต่อมาที่ hello@hotaisle.xyz ได้เลย จะไม่ส่งสแปมเด็ดขาด
    สามารถบริจาคเวลา compute สำหรับงานนี้ได้

    • เป็นบริษัทเซิร์ฟเวอร์ตัวเร่งของ AMD นี่เอง! เป็นงานที่เจ๋งมาก และหวังว่าจะมีใครสักคนรับไปทำนะ :)
    • ไม่ได้ตั้งใจจะเสียมารยาท แต่สงสัยเจตนาของข้อเสนอนี้ แค่ได้รับสิทธิ์เข้าถึงฮาร์ดแวร์ แล้วจะมีใครพอร์ตให้ฟรี ๆ หรือ? คนนั้นได้ประโยชน์อะไร?

  • FlashAttention-3 is optimized for Hopper GPUs (e.g. H100).
    FA3 จะได้ประสิทธิภาพประมาณไหนบน GPU ระดับผู้บริโภคอย่าง 3090 และ 4090?

    • เป็นของ Hopper โดยเฉพาะ จุดปรับปรุงผูกกับ ฟีเจอร์ของ Hopper อย่าง warp group และ TMA อย่างมาก
      บน 4090 อาจได้ความเร็วเพิ่มขึ้นถ้าใช้ implementation ของ FP8 attention ใน Triton: 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.
    เท่าที่เข้าใจ TMA ช่วยลดการใช้ register ก็จริง แต่สิ่งที่สำคัญกว่าคือมันทำให้ฮาร์ดแวร์จัดการ การสร้างที่อยู่ ให้ ยิ่งการคำนวณรอบข้างเร็วขึ้น การสร้างที่อยู่ก็อาจกลายเป็นคอขวดได้

  • นี่เป็นหนึ่งในการปรับปรุงที่สำคัญที่สุดใน AI ทั้งหมด มันทำให้ใช้สิ่งต่าง ๆ ได้มากขึ้นและเร็วขึ้นบนฮาร์ดแวร์เดิม และให้ ประโยชน์แทบไม่มีข้อแลกเปลี่ยน แก่ผู้ใช้ AI ส่วนใหญ่

    • ก็คงใช่สำหรับผู้ใช้ที่มี H100

  • สงสัยว่าทำไม FlashAttention เมื่อใช้ variable masking ถึงช้ากว่าตอนไม่ใช้ประมาณ 5 เท่า ถ้ารองรับ masking ได้ไม่ดี ผลของการปรับแต่งก็แทบหายไปหมด

    • ดู benchmark นั้นจากที่ไหนอยู่ครับ?

  • อยากให้ผู้เชี่ยวชาญช่วยตอบคำถามสักสองสามข้อ :)
    FlashAttention เป็นสิ่งที่ใช้แทนการคำนวณ attention ของ LLM แบบดรอปอินได้เลยหรือไม่? ใช้ได้ทุกที่ที่มีการใช้การคำนวณ “attention” หรือว่าต้องเทรน LLM แยกต่างหากให้ใช้ FA?
    FA มีความสัมพันธ์อย่างไรกับกลยุทธ์อย่าง GQA หรือ sliding window attention? เป็นแนวคิดที่ตั้งฉากกัน หรือว่าต้องมีการใช้งาน FA แยกสำหรับแต่ละกลยุทธ์?
    ล่าสุด llama.cpp เพิ่มการรองรับ FlashAttention แล้ว นั่นหมายความว่าเริ่มใช้สิ่งอย่าง CUDA kernel ที่ FlashAttention มีให้ใช่ไหม?
    สุดท้าย บทความนี้เปรียบเทียบ FlashAttention กับ Triton ไม่ใช่ว่า Triton เป็นเหมือนชั้น abstraction หรือ? ไม่สามารถทำ FA ด้วย Triton ได้หรือ? ยังไม่ค่อยเข้าใจสำนวน “FlashAttention ปะทะ Triton”

      1. เกือบถูกแล้ว เทียบเท่ากันในเชิงคณิตศาสตร์ ปัญหาในเชิงซอฟต์แวร์มีแค่เรื่องอย่างการจัดการเวอร์ชัน dependency หรือรูปแบบข้อมูลในหน่วยความจำ และ FlashAttention 2 ก็อยู่ใน HuggingFace และไลบรารียอดนิยมหลายตัวแล้ว FlashAttention 3 ก็มีแนวโน้มสูงว่าจะเข้าไปเร็ว ๆ นี้ แต่ถ้าจะรันต้องใช้ GPU H100
      2. FlashAttention 2 เพิ่มการรองรับ GQA ในอัปเดตเวอร์ชันก่อนหน้าแล้ว:
        https://github.com/Dao-AILab/flash-attention
      3. ตรงนี้เป็นการเปรียบเทียบ implementation ของ FlashAttention ตัวนี้ที่เขียนด้วย CUDA C++ ล้วน กับ implementation บน Triton ของอัลกอริทึมที่คล้ายกันซึ่งเขียนด้วย Triton: https://triton-lang.org/main/getting-started/tutorials/06-fused-attention.html
    • FlashAttention สามารถใช้แทนการคำนวณ attention ของ LLM แบบดรอปอินได้
      FlashAttention เป็นวิธีคำนวณส่วน Softmax(QK^T)V ของ attention ส่วน GQA เป็นวิธีคำนวณเมทริกซ์ Q, K, V ส่วน sliding window attention แม้จะไม่มั่นใจเท่าไร แต่น่าจะเป็นวิธีที่เปลี่ยน attention mask เพื่อควบคุมว่า query ใดมองเห็น key ใดได้บ้าง
      ผมยังไม่เคยใช้ llama.cpp แต่คำอธิบายว่าเริ่มนำ CUDA kernel มาใช้นั้นดูเหมือนจะถูกโดยรวม
      คำถามสุดท้ายหมายถึง implementation ของ FlashAttention ที่ก่อนหน้านี้เขียนด้วย Triton

  • เห็นในบทความนี้บอกว่า operation อย่าง sigmoid ช้ามาก เลยสงสัยว่า
    LLM สมัยใหม่ใช้ activation function ที่มี sigmoid หรือ Softmax อยู่มากมาย เช่น SiLU, Swish, SOLU
    ReLU มีผลเสียด้านประสิทธิภาพน้อยกว่าหรือไม่? ถ้าอย่างนั้นอาจจะดีกว่าหรือเปล่าที่จะกลับไปใช้ ReLU ดี ๆ แบบสมัยก่อน?

    • ReLU เป็นฟังก์ชันเชิงเส้นที่ถูกตัดเป็น 0 ณ จุดหนึ่งตามตัวอักษร ดังนั้นปริมาณการคำนวณจึงน้อยกว่าพวกที่มีฟังก์ชันเอ็กซ์โปเนนเชียลอยู่มาก เพียงแต่คงยากที่จะได้ผลลัพธ์ที่แข่งขันได้ด้วย activation function ที่เรียบง่ายขนาดนั้น