Lean4-TileLang Tensor Program Superoptimizer: การค้นหา GPU Kernel อัตโนมัติด้วยการอนุมานเชิงรูปแบบ เพื่อเร่งความเร็ว Attention 4.08 เท่า

ภาพรวมโครงการ

• โครงการหลัก: Lean4-TileLang Tensor Program Superoptimizer (อยู่ระหว่างการพัฒนา)
• ที่อยู่โครงการ: https://leloykun.github.io/ponder/lean4-tilelang/
• ขนาดเนื้อหา: 5,000 คำ อ่านประมาณ 18 นาที พอดแคสต์ประมาณ 25 นาที
• คำสำคัญทางเทคนิค: การอนุมานเชิงรูปแบบด้วย Lean4, การสร้าง GPU Kernel ด้วย TileLang, Tensor Program Superoptimizer, การรวมโอเปอเรเตอร์และการเพิ่มประสิทธิภาพหน่วยความจำ, การเร่งความเร็ว Attention 4.08 เท่า

导读เนื้อหา

บทความนี้มีวัตถุประสงค์เพื่อตีความแนวคิดหลักที่ Franz Louis Cesista นำเสนอในบทความบล็อกเรื่อง “Lean4-TileLang Tensor Program Superoptimizer [WIP]” อย่างเจาะลึก และร่วมกับคลังข้อมูล tile-ai/tilelang เพื่อวิเคราะห์ความสัมพันธ์ภายในและกลไกการทำงานร่วมกันระหว่าง Lean4, TileLang, กราฟการคำนวณเทนเซอร์, การเพิ่มประสิทธิภาพ GPU Kernel และ Superoptimizer อย่างครอบคลุม

บทความต้นฉบับเปิดเผยระบบนวัตกรรมที่กำลังอยู่ในขั้นตอนการพัฒนา: ระบบนี้首先ใช้ Lean4 เพื่ออธิบายกราฟการคำนวณของโครงข่ายประสาทเทียมในเชิงรูปแบบ จากนั้นจะค้นหาแผนการใช้งานที่มีประสิทธิภาพในการดำเนินการสูงกว่าแต่เทียบเท่าทางคณิตศาสตร์โดยอัตโนมัติ และสุดท้ายจะคอมไพล์แผนการเหล่านี้เป็น TileLang Kernel และดำเนินการบน GPU อย่างมีประสิทธิภาพ

บทความนี้จะพยายามใช้ภาษาที่เข้าใจง่ายที่สุด เพื่ออธิบายหลักการทำงานของระบบ ความสำคัญของการมีอยู่ของมัน รวมถึงกลยุทธ์การเพิ่มประสิทธิภาพของ Kernel ทั่วไปบางส่วนในระบบ

สารบัญบทความ

• หนึ่ง、ข้อสรุปโดยรวมก่อน
• สอง、แก่นแท้ของ TileLang
• สาม、คุณค่าของการนำ Lean4 มาใช้
• สี่、ความจำเป็นของการเพิ่มประสิทธิภาพอัตโนมัติ
• ห้า、สรุปผลการทดลองอย่างรวดเร็ว
• หก、Attention: เส้นทางการเพิ่มประสิทธิภาพที่ใกล้เคียงกับ FlashAttention
• เจ็ด、SwiGLU: การผสานรวมการคูณเมทริกซ์และฟังก์ชันกระตุ้นอย่างลึกซึ้ง
• แปด、Matmul: พื้นที่ในการเพิ่มประสิทธิภาพ GEMM พื้นฐานมีจำกัด
• เก้า、RMSNorm: การแปลงที่ถูกต้องไม่ได้หมายความว่าเร็วขึ้นเสมอไป
• สิบ、RMSNorm-MLP: ผลตอบแทนจากการรวมโอเปอเรเตอร์ปรากฏขึ้นอีกครั้ง
• สิบเอ็ด、กฎที่เปิดเผยจากการวิเคราะห์กรณีศึกษา
  • 1. แบนด์วิธหน่วยความจำสำคัญกว่าการคำนวณ
  • 1. การรวมโอเปอเรเตอร์เป็นแหล่งที่มาหลักของการปรับปรุงประสิทธิภาพ
  • 1. พื้นที่ในการเพิ่มประสิทธิภาพโอเปอเรเตอร์พื้นฐานอย่างอิสระมีจำกัด
  • 1. การค้นหาอัตโนมัติต้องอาศัยแบบจำลองต้นทุนที่แม่นยำ
• สิบสอง、คุณค่าระยะยาวของระบบนี้
• สิบสาม、ตำแหน่งในระบบนิเวศของ TileLang
• สิบสี่、การเปรียบเทียบเป็นรูปเป็นร่าง
• สรุป

หนึ่ง、ข้อสรุปโดยรวมก่อน

จากเนื้อหาของ “Lean4-TileLang Tensor Program Superoptimizer” สามารถสรุประบบทั้งหมดได้ดังนี้:

ใช้ Lean4 เพื่อกำหนด “สิ่งที่ต้องคำนวณ” ในเชิงรูปแบบ 借助 Superoptimizer เพื่อค้นหา “วิธีการเขียนใหม่ที่เทียบเท่า” โดยอัตโนมัติ และสุดท้ายผ่าน TileLang เพื่อสร้าง “Kernel ที่สามารถทำงานบน GPU ได้อย่างมีประสิทธิภาพ”

กล่าวอีกนัยหนึ่ง นี่ไม่ใช่แค่การเขียน Kernel ประสิทธิภาพสูงด้วยมือสองสามตัว แต่เป็นการพยายามสร้างห่วงโซ่การเพิ่มประสิทธิภาพอัตโนมัติ:

คำจำกัดความทางคณิตศาสตร์ของโอเปอเรเตอร์โครงข่ายประสาทเทียม
↓
กราฟการคำนวณ Lean4
↓
การแปลงที่เทียบเท่าและการค้นหา
↓
สร้าง TileLang Kernel
↓
คอมไพล์แบบ JIT
↓
ทำงานบน GPU และทดสอบสมรรถนะ

ในกระบวนการนี้ หน้าที่ของแต่ละองค์ประกอบมีดังนี้:

โดยสรุป: Lean4 รับประกันว่า “การคำนวณถูกต้อง” ในขณะที่ TileLang รับผิดชอบ “การทำงานมีประสิทธิภาพ”

สอง、แก่นแท้ของ TileLang

tile-ai/tilelang เป็นภาษาเฉพาะโดเมน (DSL) ที่ออกแบบมาโดยเฉพาะสำหรับการเขียน Kernel GPU/CPU ประสิทธิภาพสูง

เป้าหมายการออกแบบคือการช่วยให้นักพัฒนาสามารถเขียน Kernel ที่มีประสิทธิภาพใกล้เคียงกับระบบระดับล่างอย่าง CUDA, CUTLASS, TVM โดยใช้โค้ดสไตล์ Python ที่ค่อนข้างกระชับ

ในคลังข้อมูล tile-ai/tilelang โอเปอเรเตอร์ทั่วไปที่ TileLang รองรับ ได้แก่:

• GEMM / Matmul
• Dequant GEMM
• FlashAttention
• Linear Attention
• Convolution
• DeepSeek MLA
• Kernel AI ที่ซับซ้อนอื่น ๆ

โค้ด TileLang ดูคล้ายกับ Python แต่ไม่ใช่ Python ทั่วไป มันจะถูกแปลงเป็น Kernel ระดับล่างผ่านการคอมไพล์แบบทันที (JIT) รูปแบบการเขียนทั่วไปและคำอธิบาย通俗มีดังนี้:

สามารถเข้าใจ TileLang ได้ดังนี้:

ภาษาการเขียน Kernel ที่ต่ำกว่า PyTorch แต่เขียนง่ายกว่า CUDA

PyTorch ถนัดในการแสดงตรรกะของโมเดล ในขณะที่ TileLang ถนัดในการแสดง “โอเปอเรเตอร์นี้ควรแบ่งบล็อกบน GPU อย่างไร, ย้ายข้อมูลอย่างไร, ใช้หน่วยความจำที่ใช้ร่วมกันอย่างไร, และเรียกใช้ Tensor Core อย่างไร”

สาม、ทำไมต้องใช้ Lean4?

เนื่องจาก TileLang สามารถเขียน Kernel ประสิทธิภาพสูงได้แล้ว เหตุใดจึงต้องนำ Lean4 มาใช้?

ประเด็นสำคัญอยู่ที่: การเพิ่มประสิทธิภาพ GPU Kernel มักเกี่ยวข้องกับการแปลงที่เทียบเท่าทางคณิตศาสตร์

ในการใช้งาน Attention มาตรฐาน โค้ดมักจะมีลักษณะดังนี้:

scores = Q @ Kᵀ
probs = softmax(scores)
out = probs @ V

อย่างไรก็ตาม การใช้งาน FlashAttention ประสิทธิภาพสูงจะไม่เขียนเมทริกซ์ scores ทั้งหมดลงในหน่วยความจำ GPU แต่จะใช้กลยุทธ์การคำนวณแบบแบ่งบล็อก และอาศัยเทคนิค softmax ออนไลน์เพื่อรักษาสถานะตัวกลาง

ที่นี่ต้องรับประกันประเด็นสำคัญ:

ผลลัพธ์ของการคำนวณออนไลน์แบบแบ่งบล็อก

เทียบเท่าทางคณิตศาสตร์อย่างสมบูรณ์กับผลลัพธ์ของการคำนวณต่อไปนี้:

scores เต็ม + softmax + matmul

อีกตัวอย่างหนึ่งคือการแปลงใน RMSNorm:

out = (X * S) @ W

ภายใต้เงื่อนไขเฉพาะ มันสามารถเขียนใหม่เป็น:

out = (X @ W) * S

โดยที่ S คือปัจจัยสเกลที่กระจายตามแถว หากเงื่อนไขไม่เป็นไปตามนี้ การแปลงจะผิดพลาด

บทบาทของ Lean4 คือการให้พื้นฐานที่เข้มงวดยิ่งขึ้นสำหรับกราฟการคำนวณและการแปลงเหล่านี้

มันสามารถช่วยให้ระบบอธิบายสิ่งต่อไปนี้ได้อย่างชัดเจน:

• รูปร่าง (shape) ของอินพุตและเอาต์พุต
• ข้อจำกัดของแกน (axis)
• ความหมายของการลดรูป (reduction)
• กฎของการกระจาย (broadcast)
• แกนหดตัว (contraction axis) ของการคูณเมทริกซ์
• รายละเอียดของการดำเนินการแบบทีละองค์ประกอบ
• การแปลงใดที่ถูกต้องตามกฎหมาย
• กราฟการคำนวณใดที่เทียบเท่ากันในเชิงตรรกะ

ดังนั้น ในระบบนี้ Lean4 ทำหน้าที่เป็น “ชั้นคณิตศาสตร์และการอนุมาน” ในขณะที่ TileLang รับผิดชอบ “ชั้นการดำเนินการประสิทธิภาพสูง”

สี่、เหตุใดจึงต้องมีการเพิ่มประสิทธิภาพอัตโนมัติเช่นนี้?

ปัจจุบัน โมเดลโครงข่ายประสาทเทียมมีความซับซ้อนมากขึ้นเรื่อย ๆ โครงสร้างใหม่ ๆ เกิดขึ้นอย่างต่อเนื่อง เช่น ตัวแปร Attention ใหม่, โครงสร้าง MLP, วิธีการ Normalization และกลไก Gating ฯลฯ

หากทุกครั้งที่มีโครงสร้างใหม่เกิดขึ้น จำเป็นต้องเขียน CUDA kernel ด้วยมือชุดหนึ่ง จะต้องเผชิญกับปัญหาหลายประการ:

1. ต้นทุนการพัฒนาสูงมาก
2. ต้องมีประสบการณ์การเขียน GPU kernel ที่เชี่ยวชาญเป็นอย่างมาก
3. มีแนวโน้มที่จะเกิดข้อผิดพลาดสูง
4. ยากที่จะเปรียบเทียบโครงสร้างตัวเลือกจำนวนมากได้อย่างรวดเร็ว
5. ผลการทดลองขนาดเล็กอาจไม่สะท้อนประสิทธิภาพจริงในการฝึกขนาดใหญ่

ดังที่ผู้เขียนกล่าวไว้ในตอนต้นของบล็อกต้นฉบับ เขาไม่เห็นด้วยกับการวิจัยอัตโนมัติแบบ “fuzzing โค้ด” ในขนาดเล็กเท่านั้น เนื่องจากการทำงานที่ดีในขนาดเล็กไม่ได้หมายความว่าจะทำงานได้ดีในขนาดใหญ่เช่นกัน เพื่อเปรียบเทียบสถาปัตยกรรมที่แตกต่างกันอย่างยุติธรรม อย่างน้อยต้องพิจารณาพร้อมกัน:

• GPU / TPU kernel ได้รับการปรับให้เหมาะสมเพียงพอหรือไม่
• optimizer ถูกเลือกอย่างเหมาะสมหรือไม่
• วิธีการกำหนดพารามิเตอร์เหมาะสมหรือไม่
• ไฮเปอร์พารามิเตอร์เปลี่ยนแปลงอย่างไรตามความกว้าง ความลึก ขนาดแบตช์ และจำนวนขั้นตอนการฝึก
• กฎขนาดเล็กสามารถถ่ายโอนไปยังสถานการณ์ขนาดใหญ่ได้หรือไม่

ดังนั้น การสร้าง kernel ประสิทธิภาพสูงโดยอัตโนมัติเป็นเพียงส่วนหนึ่งของเป้าหมายที่ใหญ่กว่า:

กำหนดโครงสร้างโครงข่ายประสาทเทียมด้วยระบบเชิงรูปแบบ
        ↓
สร้าง kernel ประสิทธิภาพสูงโดยอัตโนมัติ
        ↓
ปรับ optimizer / parametrization โดยอัตโนมัติ
        ↓
ศึกษา scaling law
        ↓
เปรียบเทียบสถาปัตยกรรมที่แตกต่างกันอย่างยุติธรรมยิ่งขึ้น

บทความนี้เน้นการตีความขั้นตอน “Lean4 → TileLang kernel” เป็นหลัก

ห้า、สรุปผลการทดลองอย่างรวดเร็ว

จากข้อมูล benchmark ที่ให้ไว้ในบทความต้นฉบับ ผลลัพธ์ของ workload หลายรายการมีดังนี้:

จากนี้สามารถสรุปข้อสังเกต直观ได้หลายประการ:

1. การเร่งความเร็วของ Attention มีนัยสำคัญมาก ประมาณ 4 เท่า
2. SwiGLU ก็มีการปรับปรุงที่ชัดเจน ประมาณ 2 เท่า
3. การปรับปรุงของ Matmul ค่อนข้างจำกัด เนื่องจากการคูณเมทริกซ์ได้รับการปรับให้เหมาะสมอย่างมากอยู่แล้ว
4. RMSNorm เมื่อทดสอบเดี่ยว ๆ กลับช้าลง
5. แต่เมื่อรวม RMSNorm กับ MLP แล้ว ประสิทธิภาพก็ดีขึ้น

สิ่งนี้เปิดเผยข้อเท็จจริงสำคัญ:

ข้อได้เปรียบของ TileLang และตัวเพิ่มประสิทธิภาพอัตโนมัตินั้น อยู่ที่ “การรวมโอเปอเรเตอร์” และ “การลดการอ่านเขียนหน่วยความจำ GPU” มากกว่าแค่การทำให้โอเปอเรเตอร์พื้นฐานตัวเดียวทำงานเร็วขึ้น

หก、Attention: แนวคิดการเพิ่มประสิทธิภาพที่ใกล้เคียงกับ FlashAttention

รูปแบบทางคณิตศาสตร์ของ Attention มีดังนี้:

scores = Q @ Kᵀ
probs = softmax(scores)
out = probs @ V

คำอธิบาย Lean4 ที่สอดคล้องกันโดยประมาณคือ:

let q ← input “Q” shapeHTqDh
let k ← input “K” shapeHTkvDh
let v ← input “V” shapeHTkvDh
let scores ← matmul q k “d_h” “d_h”
let probs ← softmaxSubdag scores “t_kv”
let out ← matmul probs v “t_kv” “t_kv”
output out

นี่คือกราฟการคำนวณ Attention มาตรฐาน

ปัญหาคือ ขนาดของ scores จะใหญ่มาก จากการกำหนดค่าในบทความต้นฉบับ:

h = 16
t_q = 4096
t_kv = 4096
d_h = 128

ดังนั้นรูปร่างของ scores จะเป็น:

16 × 4096 × 4096

หากสร้างเมทริกซ์นี้เต็มรูปแบบ แล้วเขียนลงหน่วยความจำ GPU จากนั้นอ่านกลับมาทำ softmax และสุดท้ายคูณกับ V ปริมาณการรับส่งข้อมูลหน่วยความจำ GPU จะมหาศาล

Derivation of top-1 kernel

กุญแจสำคัญในการใช้งานประสิทธิภาพสูงคือการหลีกเลี่ยงการเขียนเมทริกซ์ scores ทั้งหมดลงในหน่วยความจำ GPU

TileLang kernel ที่สร้างขึ้นใช้วิธีการคำนวณออนไลน์แบบแบ่งบล็อกคล้ายกับ FlashAttention:

固定一块 Q
循环读取 K/V 的小块
每次计算局部 QKᵀ
做局部 softmax 统计
更新全局 softmax 状态
累加输出

ในโค้ด TileLang มีตัวแปรสถานะสำคัญหลายตัว:

state_pass0_m
state_pass0_l
state_pass0_o

ความหมายของตัวแปรเหล่านี้สามารถเข้าใจได้ดังนี้:

เหตุใดจึงต้องมีสถานะเหล่านี้?

因为 softmax 不能简单地分块各自计算。为了保证分块计算与完整 softmax 的等价性，必须维护全局的最大值和归一化因子。

กระบวนการ softmax ออนไลน์นี้สามารถเข้าใจได้ดังนี้:

ทุกครั้งที่มี K/V บล็อกใหม่เข้ามา
ก็จะรวมมันเข้ากับสถิติ softmax ปัจจุบัน
พร้อมกันนั้นก็แก้ไขผลลัพธ์สะสมเก่า

ด้วยวิธีนี้จึงไม่จำเป็นต้องเก็บเมทริกซ์ scores ทั้งหมด

นี่คือแนวคิดหลักของอัลกอริทึมตระกูล FlashAttention:

ใช้การคำนวณบนชิปมากขึ้น เพื่อแลกกับการอ่านเขียนหน่วยความจำ GPU น้อยลง

ดังนั้น สำหรับสถานการณ์ Attention นี้ เราจึงได้รับการปรับปรุงประสิทธิภาพประมาณ 4 เท่า และบทความต้นฉบับระบุเป็นพิเศษว่าผลลัพธ์นี้เทียบได้กับ Flash Attention 2

เจ็ด、SwiGLU: การหลอมรวมการคูณเมทริกซ์สองครั้งกับฟังก์ชันกระตุ้น

SwiGLU เป็นโครงสร้างที่พบได้บ่อยมากใน MLP ของโมเดลภาษาขนาดใหญ่ สามารถลดรูปเป็นขั้นตอนการคำนวณต่อไปนี้:

up = X @ W_up  
gate = X @ W_gate  
out = up * silu(gate)  

โดยที่ฟังก์ชัน silu ถูกกำหนดเป็น:

silu(x) = x * sigmoid(x)  

ในคำอธิบายกราฟการคำนวณ Lean4 มันถูกแสดงเป็น:

let x ← input "X" shapeMD  
let wUp ← input "W_up" shapeDN  
let wGate ← input "W_gate" shapeDN  
let gatePre ← matmul x wGate "d" "d"  
let up ← matmul x wUp "d" "d"  
let gate ← silu gatePre  
let out ← mul gate up  
output out  

หากใช้เฟรมเวิร์กทั่วไปในการดำเนินการ อาจสร้างเทนเซอร์ตัวกลางหลายตัว เช่น:

gatePre  
up  
gate  
out  

เทนเซอร์ตัวกลางเหล่านี้แต่ละครั้งจำเป็นต้องเขียนกลับไปยังหน่วยความจำ GPU แล้วอ่านกลับมาสำหรับขั้นตอนถัดไป

การ推导 Kernel ที่ดีที่สุด

วิธีการของ TileLang kernel คือการรวมขั้นตอนทั้งหมดนี้เข้าด้วยกัน:

1. คำนวณ X @ W_up ใน kernel เดียวกัน
2. คำนวณ X @ W_gate ใน kernel เดียวกัน
3. ดำเนินการ SiLU กับ gate โดยตรง
4. คูณกับ up โดยตรง
5. สุดท้าย เขียนเฉพาะผลลัพธ์สุดท้ายกลับไปยังหน่วยความจำ GPU

รูปแบบโค้ดหลักคล้ายกับ:

T.gemm(input_0, input_4, matmul_8, clear_accum=False)    # up  
T.gemm(input_0, input_10, matmul_14, clear_accum=False)  # gate  

จากนั้น รวมการดำเนินการ SiLU และการคูณ:

matmul_14[i0, i1] = matmul_14[i0, i1] / (1.0 + T.exp2((-(matmul_14[i0, i1])) * scale))  
matmul_14[i0, i1] = matmul_8[i0, i1] * matmul_14[i0, i1]  

นอกจากนี้ยังใช้เทคนิค fast math:

scale = 1.44269504  # log2(e)

原因在于:

exp(x) = exp2(x * log2(e))

บน GPU exp2 มักจะคำนวณได้อย่างมีประสิทธิภาพมากกว่า

การปรับปรุงประสิทธิภาพของ SwiGLU มาจากสองด้านหลัก:

1. การคูณเมทริกซ์สองครั้ง (matmul) ใช้ tile อินพุตเดียวกัน
2. การดำเนินการ SiLU และการคูณไม่จำเป็นต้องเขียนผลลัพธ์ตัวกลางลงในหน่วยความจำ GPU แยกต่างหาก

ดังนั้น สถานการณ์นี้จึงได้รับการเร่งความเร็วประมาณ 2.14 เท่า

แปด、Matmul: พื้นที่ในการเพิ่มประสิทธิภาพ GEMM พื้นฐานค่อนข้างจำกัด

การคูณเมทริกซ์ (Matmul) เป็นการดำเนินการพื้นฐานที่สุด:

O = X @ W

กราฟการคำนวณ Lean4 นั้น简洁มาก:

let x ← input "X" shapeMD  
let w ← input "W" shapeDN  
let out ← matmul x w "d" "d"  
output out  

การ推导 Kernel ที่ดีที่สุด

TileLang kernel ที่สร้างขึ้นคือ GEMM แบบแบ่งบล็อก (tiled) มาตรฐาน:

把 X 的一块搬到 shared memory  
把 W 的一块搬到 shared memory  
用 T.gemm 做 tile-level 矩阵乘  
沿 d 维循环累加  
最后写回 O  

โค้ดหลักคล้ายกับ:

for k_pass0 in T.Pipelined(T.ceildiv(d, block_d), num_stages=num_stages):  
T.copy(X[gy * block_m, k_pass0 * block_d], input_0)  
T.copy(W[k_pass0 * block_d, gx * block_n], input_4)  
T.gemm(input_0, input_4, matmul_8, clear_accum=False)  

มันใช้เทคนิคการเพิ่มประสิทธิภาพ GEMM ทั่วไปของ GPU รวมถึง:

• การแบ่งบล็อก (tiling)
• หน่วยความจำที่ใช้ร่วมกัน (shared memory)
• การสะสมแบบ fragment (fragment accumulation)
• ลูปแบบไปป์ไลน์ (pipelined loop)
• การสลับ (swizzle)
• Tensor Core GEMM

อย่างไรก็ตาม การคูณเมทริกซ์เป็นหนึ่งในโอเปอเรเตอร์ที่ได้รับการปรับให้เหมาะสมที่สุดบน GPU อยู่แล้ว ไลบรารีอย่าง cuBLAS, CUTLASS, แบ็กเอนด์ของ PyTorch และ torch.compile ต่างก็ดำเนินการเพิ่มประสิทธิภาพอย่างมากมาแล้ว

ดังนั้น เวอร์ชัน TileLang จึงเร็วกว่า torch.compile เพียงประมาณ 1.29 เท่า ซึ่งไม่น่าแปลกใจ

จากนี้สามารถสรุปข้อสรุปสำคัญได้:

คุณค่าสูงสุดของการสร้าง TileLang kernel โดยอัตโนมัติ ไม่จำเป็นต้องอยู่ที่การเอาชนะไลบรารี GEMM ที่เติบโตเต็มที่เพียงลำพัง แต่อยู่ที่ความสามารถในการรวม GEMM เข้ากับการดำเนินการอื่น ๆ อย่างมีประสิทธิภาพ

เก้า、RMSNorm: การแปลงที่ถูกต้องไม่ได้หมายความว่าเร็วขึ้นเสมอไป

สูตรอย่างง่ายของ RMSNorm มีดังนี้:

S = 1 / sqrt(mean(X²) + eps)  
norm = X * S  
out = norm @ W  

ใน Lean4 ก่อนอื่นจะอธิบายกราฟย่อยของ RMSNorm:

let x2 ← square x  
let mean ← meanAlongSubdag x2 axis  
let eps ← epsConst  
let denom ← add mean eps  
let scale ← rsqrt denom  
mul x scale  

ตามด้วยการคูณเมทริกซ์:

let out ← matmul norm w "d" "d"  

การใช้งาน TileLang ที่สร้างขึ้นโดยอัตโนมัติจะแบ่งออกเป็นสอง Kernel

Kernel แรก负责คำนวณค่า scale สำหรับแต่ละแถว:

S[gy * block_m + i0] = 1.0 / T.sqrt(row_state[i0] / d + eps)  

Kernel ที่สองดำเนินการคูณเมทริกซ์ และคูณด้วย scale ในตอนท้าย:

T.gemm(X_tile, W_tile, acc_0, clear_accum=False)  

for i0, i1 in T.Parallel(block_m, block_n):  
acc_0[i0, i1] = acc_0[i0, i1] * S[gy * block_m + i0]  

ที่นี่ใช้การแปลงทางพีชคณิต:

(X * S) @ W = (X @ W) * S  

ตราบใดที่ S เป็นปัจจัยสเกลที่กระจายตามแถว การแปลงนี้ก็สมเหตุสมผล

อย่างไรก็ตาม กลับเกิดปัญหาด้านประสิทธิภาพ

เนื่องจากการใช้งานแบบนี้จำเป็นต้อง…

การเขียนใหม่เชิงลึกและเวอร์ชันลดขนาด

หลังจากดำเนินการ Kernel หนึ่งเพื่อคำนวณ S แล้ว จำเป็นต้องเขียนผลลัพธ์กลับไปยังหน่วยความจำ GPU จากนั้น ต้องเริ่ม Kernel ใหม่เพื่อดำเนินการคูณเมทริกซ์ (matmul) และอ่านข้อมูล S ที่เก็บไว้ก่อนหน้านี้จากหน่วยความจำ GPU อีกครั้ง สุดท้าย คูณค่า S ที่อ่านได้นี้เข้ากับผลลัพธ์เอาต์พุต

สำหรับโอเปอเรเตอร์ที่มีโครงสร้างค่อนข้างเรียบง่ายอย่าง RMSNorm การเพิ่มการเริ่ม Kernel (kernel launch) และการดำเนินการอ่านเขียนข้อมูลตัวกลางอีกครั้ง อาจลบล้างการปรับปรุงประสิทธิภาพที่ได้จากการเพิ่มประสิทธิภาพได้อย่างสมบูรณ์

ผลการทดสอบสุดท้ายก็ยืนยันประเด็นนี้:

TileLang: 0.206 ms  
torch.compile: 0.174 ms  
Speedup: 0.842x  

กระบวนการ推导 Kernel ที่ดีที่สุด

นั่นหมายความว่า ในกรณีเฉพาะนี้ เวอร์ชันที่สร้างโดย TileLang กลับช้ากว่า torch.compile

สิ่งนี้เปิดเผยประเด็นสำคัญ:

Superoptimizer แม้จะสามารถค้นหาแผนการใช้งานที่เทียบเท่าทางคณิตศาสตร์ได้อย่างสมบูรณ์ แต่ความเทียบเท่าทางคณิตศาสตร์ไม่ได้หมายความว่าจะมีประสิทธิภาพในการทำงานจริงที่เหนือกว่า

ประสิทธิภาพจริงของ GPU ได้รับผลกระทบจากปัจจัยที่ซับซ้อนหลายประการร่วมกัน รวมถึง:

• ค่าใช้จ่ายในการเริ่ม (launch) Kernel
• แบนด์วิธหน่วยความจำ GPU และความหน่วงในการอ่านเขียน
• ประสิทธิภาพการใช้หน่วยความจำที่ใช้ร่วมกัน (shared memory)
• แรงกดดันของทรัพยากรรีจิสเตอร์ (register)
• ระดับการครอบครอง (occupancy)
• อัตราการ命中แคช (cache)
• คุณภาพของการเพิ่มประสิทธิภาพของคอมไพเลอร์ระดับล่าง

สิบ、RMSNorm-MLP: ผลตอบแทนจากการรวมโอเปอเรเตอร์ปรากฏขึ้นอีกครั้ง

RMSNorm-MLP เป็นโอเปอเรเตอร์แบบรวมที่ใกล้เคียงกับโครงสร้างโมเดลจริงมากขึ้น โดยมีขั้นตอนการคำนวณดังนี้:

norm = rmsnorm(X)  
up = norm @ W_up  
gate = norm @ W_gate  
out = up * gate  

เมื่อแสดงด้วยกราฟ Lean4 สามารถเข้าใจได้ดังนี้:

let norm ← rmsNormSubdag x "d"  
let up ← matmul norm wUp "d" "d"  
let gate ← matmul norm wGate "d" "d"  
let out ← mul up gate  
output out  

ในการใช้งาน TileLang ยังคงต้องคำนวณค่า S สำหรับแต่ละแถวก่อน:

S = scale_kernel(values['X'], eps_value)


⚠️ หมายเหตุ: เนื้อหาได้รับการแปลโดย AI และตรวจสอบโดยมนุษย์ หากมีข้อผิดพลาดโปรดแจ้ง

    

        
☕ สนับสนุนค่ากาแฟทีมงาน
        
หากคุณชอบบทความนี้ สามารถสนับสนุนเราได้ผ่าน PromptPay
        

            
        
        

            SCAN TO PAY WITH ANY BANK
        
    
                                                        
本文来自网络搜集，不代表คลื่นสร้างอนาคต立场，如有侵权，联系删除。转载请注明出处：https://www.itsolotime.com/th/archives/34713