Zyphra เปิดตัว ZYFA เร่งความเร็ว FP8 บน AMD MI355X สูงสุด 1.76 เท่า

Zyphra ได้เปิดตัว ZYFA (Zyphra FlashAttention on AMD) ซึ่งเป็นชุด prefill attention kernels สำหรับ FP8 และ BF16 บน GPU รุ่น AMD Instinct MI355X โดย ZYFA คือเอนจิน attention แบบ FP8 ของ Zyphra สำหรับการทำ long-context prefill และเปิดให้ใช้งานแบบ open source ภายใต้ลิขสิทธิ์ Apache 2.0

จากการทำงานบน activation ของโมเดลจริง FP8 kernels ของ ZYFA สามารถทำงานได้เร็วกว่า AITER ซึ่งเป็นไลบรารี kernel มาตรฐานของ AMD ถึง 1.16-1.76 เท่า และที่ head dimension 128 พบว่า ZYFA มีค่าความผิดพลาดของ FP8 ต่ำกว่า AITER ถึง 62-78%

TL;DR

• Framework: ไลบรารี kernel สำหรับ Causal forward (prefill) attention ในรูปแบบ FP8 e4m3 และ BF16 รองรับ head dim 128 และ MLA shape (192/128)
• Runs on: รองรับเฉพาะ AMD Instinct MI355X (CDNA4, gfx950) บน ROCm 7.2.3 และ PyTorch 2.12 โดยมีอินเทอร์เฟซแบบ flash_attn_func ที่นำไปใช้แทนที่ได้ทันที
• Performance: ให้ throughput สำหรับ FP8 สูงกว่า AITER 1.16-1.76 เท่าบน activation จริง ส่วน BF16 มีประสิทธิภาพใกล้เคียงกันที่ 1.01-1.02 เท่า
• Best: ทำได้ 2,672 TFLOP/s สำหรับ FP8 ที่ 32K tokens หรือคิดเป็น 53% ของ FP8 datasheet peak (ขณะที่ AITER ทำได้ 39%)
• Worst: เร็วขึ้นเพียง 1.16 เท่าบน activation ของ Mistral-7B เนื่องจากความเร็วนาฬิกาของ FP8 ลดลงเหลือ 1,628 MHz เมื่อเทียบกับ AITER ที่ 2,108 MHz

ZYFA คืออะไร?

ZYFA เผยแพร่ในรูปแบบรายงานทางเทคนิคพร้อมซอร์สโค้ดจาก Zyphra โดยเป็นส่วนขับเคลื่อนหลักของ Zyphra Optimization Engine (ZOE) ที่ใช้งานบน Zyphra Cloud สำหรับ BF16 kernels นั้นเขียนด้วยภาษา HIP ส่วน FP8 kernels เป็นแบบ hand-scheduled assembly ที่สร้างโดย Python generator

ทำไม FP8 attention ถึงมีประสิทธิภาพต่ำบน MI355X?

แม้คำสั่ง matrix สำหรับ FP8 จะทำงานได้เร็วเป็น 2 เท่าของจุดสูงสุดใน BF16 แต่ FP8 ช่วยเร่งความเร็วได้เฉพาะส่วน matrix product (QKᵀ และ PV) เท่านั้น ขณะที่การคำนวณ softmax ที่อยู่ตรงกลางยังคงทำงานบนหน่วย vector ด้วยความเร็วเท่าเดิมในทุกระดับความแม่นยำ

สำหรับการคิวรี 32 ชุดเทียบกับคีย์ 128 ชุด การทำงานของ matrix ในรูปแบบ FP8 ใช้เวลา 1,024 cycles ขณะที่ FP8 softmax แบบแม่นยำใช้เวลาถึง 896 vector cycles ทำให้ kernel ต้องการการทำงานที่ซ้อนทับกัน (overlap) อย่างสมบูรณ์เพื่อรักษาประสิทธิภาพ ผลตรวจสอบ thread traces ของ Zyphra พบว่าท่อส่งข้อมูล matrix ของ AITER ว่างงานถึง 31.7% ของรอบการทำงาน เนื่องจาก MI355X ไม่ได้ติดตามการเสร็จสิ้นของคำสั่ง matrix ในฮาร์ดแวร์

ZYFA ทำงานอย่างไร?

ZYFA ใช้ 3 เทคนิคหลักเพื่อแก้ปัญหาดังกล่าว:

1. Lean P: เปลี่ยนการใช้เลขยกกำลังเป็นการคูณ-บวกเพียงครั้งเดียว โดยอาศัยคุณสมบัติของรูปแบบบิต e4m3 ช่วยลด softmax vector cycles ต่อบล็อกจาก 896 เหลือเพียง 384

2. Two-wave ping-pong schedule: ใช้ 2 wave ร่วมกันในแต่ละ SIMD โดยตัวหนึ่งส่งคำสั่ง matrix ขณะที่อีกตัวคำนวณความน่าจะเป็นแล้วสลับกัน เทคนิคนี้ช่วยให้ท่อส่งข้อมูล matrix ทำงานได้ถึง 86% ในสภาวะคงที่

3. Per-block scaling และ rounded row sum: บล็อกความน่าจะเป็นจะถูกปรับสัดส่วนด้วยเลขยกกำลังของ 2 ก่อนปัดเศษ โดยคำสั่ง scaled matrix ของ CDNA4 จะช่วยลดภาระงานลงได้โดยไม่มีต้นทุนเพิ่ม

ZYFA เร็วแค่ไหน?

การเร่งความเร็ว FP8 บน activation จริงที่ 16K พบว่าเร็วกว่าเดิม 1.16 เท่าใน Mistral-7B, 1.17 เท่าใน Qwen3-8B และสูงถึง 1.76 เท่าใน DeepSeek-V2-Lite ส่วนผลทดสอบ Gaussian benchmarks ที่ 32K ตัว head-dim-128 FP8 kernel ทำความเร็วได้ 2,672 TFLOP/s และในรูปทรงแบบ MLA ทำได้ถึง 2,464 TFLOP/s ซึ่งเร็วกว่า AITER เกือบ 2 เท่า

ในด้านประสิทธิภาพพลังงาน ZYFA ให้ประสิทธิภาพ 1.94 และ 1.91 TFLOP/J บนเลเยอร์ของ Mistral-7B และ Qwen3-8B ตามลำดับ ขณะที่ AITER ทำได้ 1.60 และ 1.56 TFLOP/J

ทำไมประสิทธิภาพในโมเดลจริงถึงน้อยกว่า benchmark?

ความเร็วนาฬิกาของ MI355X แปรผันตามจำนวนบิตที่มีการพลิก (flip) โดย AITER ปล่อยให้ความน่าจะเป็นค่าน้อยกลายเป็นศูนย์ ทำให้มีการสลับบิตน้อยและรักษาความเร็วนาฬิกาได้สูงกว่า ในเลเยอร์ Mistral-7B นั้น ZYFA ทำงานที่ 1,628 MHz ขณะที่ AITER อยู่ที่ 2,108 MHz อย่างไรก็ตาม ZYFA ยังคงใช้จำนวนรอบการทำงาน (cycles) รวมน้อยกว่า

FP8 attention ของ ZYFA มีความแม่นยำแค่ไหน?

เมื่อเทียบกับ fp64 พบว่าค่าความผิดพลาดของ ZYFA ต่ำกว่า AITER อย่างเห็นได้ชัด โดยคิดเป็นเพียง 0.22 เท่าใน Mistral-7B และ 0.38 เท่าใน Qwen3-8B นอกจากนี้ผลลัพธ์ของ ZYFA ยังมีความเสถียรระดับ bitwise identical จากการรันต่อเนื่อง 7,200 ครั้ง

ZYFA เปรียบเทียบกับ attention kernels อื่นๆ อย่างไร?

ประเด็นสำคัญ

• FP8 attention บน MI355X มีคอขวดอยู่ที่ softmax ไม่ใช่ matrix
• เทคนิค Lean P ช่วยเปลี่ยนงานหนักให้กลายเป็นการคูณ-บวกแบบง่าย
• ค่าความผิดพลาดของ FP8 ลดลง 62-78% เมื่อเทียบกับมาตรฐานเดิม
• ปัจจุบันรองรับเฉพาะ Prefill เท่านั้น ยังไม่รองรับ decode หรือ paged KV และใช้งานได้เฉพาะรุ่น MI355X