การเขียน GPU kernels ด้วยมือต้องอาศัยความอดทน ความพยายาม และก็สติของคุณด้วย ผมได้ทดลองใช้ LLM สมัยใหม่ (Claude Opus 5, GPT 5.6 Sol) เพื่อให้มันพิมพ์ kernels ออกมาให้ผมเลย และผมจะพูดถึงสิ่งที่ได้เรียนรู้จากการทำแบบนั้น
พูดตรงๆ การเขียน GPU kernels เป็นงานที่ตรวจสอบยืนยันได้ง่ายอย่างน่าอาย
อย่างแรก คุณตัดสินใจว่าจะเขียน kernel สำหรับ operation ใดบ้าง อย่างที่สอง คุณเขียนเวอร์ชันแรกและทำให้แน่ใจว่ามัน compile ผ่านโดยไม่มี error ที่ชัดเจน อย่างที่สาม คุณตรวจสอบความถูกต้องด้วย reference implementation ที่ช้า ถ้ามันไม่ตรงกัน คุณก็พยายามแก้ไขปัญหาความถูกต้อง พอทำเสร็จแล้ว คุณก็ benchmark เวลาการทำงานของ kernel เวอร์ชันต่อๆ ไปก็ต่อยอดจากจุดนี้ และคุณก็ปรับให้เหมาะสมเรื่อยๆ จนกว่าจะถึง roofline metrics หรือพอใจกับผลลัพธ์
ลูปที่ตรวจสอบยืนยันได้: การพัฒนา GPU kernel
ภาพด้านบนแสดงมันเป็นลูปที่ตรวจสอบยืนยันได้:
- การตรวจสอบการ compile คือลูปเล็กๆ ที่รวดเร็ว (B <-> C) ก่อนที่คุณจะคิดถึงเรื่องความถูกต้องด้วยซ้ำ
- การตรวจสอบความถูกต้อง (D <-> E <-> F) คือแกนหลักของลูปที่ให้ผลตอบแทนจากการตรวจสอบยืนยัน นี่คือส่วนที่ทำให้การเขียน kernel เป็นโจทย์ที่ดีสำหรับการตรวจสอบยืนยันอัตโนมัติ เพราะคุณมี ground-truth reference ไว้ใช้ตรวจเทียบ
- การปรับให้เหมาะสม (G -> H -> กลับไป D) ใช้ลูปความถูกต้องเดียวกันนี้ซ้ำสำหรับทุกเวอร์ชันใหม่ เพราะ kernel ที่เร็วแต่ผิดนั้นไร้ค่า ความถูกต้องจึงต้องถูกตรวจสอบในทุกเวอร์ชัน
- การตรวจสอบ roofline/ความพอใจ (I) คือลูปวงนอกที่ตัดสินใจว่าจะปรับให้เหมาะสมต่อหรือหยุด
เวอร์ชันแรก
ภาพด้านบนยังเป็นมุมมองระดับสูง และปีศาจมันซ่อนอยู่ในรายละเอียด เราจำเป็นต้องทำให้แน่ใจว่า LLM agent ของเรามีบริบทที่จำเป็นทั้งหมดเพื่อเริ่มเขียนเวอร์ชันแรกที่ดีได้
CUDA DSLs เริ่มมีบทบาทตรงนี้ Triton, CuTeDSL และ Tilelang เป็นตัวที่เริ่มต้นได้ง่ายมากใน Python เส้นโค้งการเรียนรู้ก็ชันน้อยกว่า CUDA C++ แต่ abstraction ใน DSL เหล่านั้นอาจทำให้ agent ของเรางงยิ่งกว่าเดิม เราจำเป็นต้องมีวิธีส่งบริบทของ abstraction เหล่านั้นให้กับ agent
LLM สมัยใหม่รู้อยู่แล้วว่าจะเขียน Triton ที่ "ดี" อย่างไร พวกมันทำงานกับ abstraction ของ Triton ได้ดีแม้ไม่มีบริบทใดๆ แต่สำหรับ DSL อื่นๆ อย่าง CuTeDSL (ซึ่งให้การควบคุมมากกว่า Triton มาก) ผมพบว่า การมี context directory ที่ agent สามารถค้นหาความเข้าใจเกี่ยวกับ abstraction ของ DSL ได้ช่วยได้มาก
ตัวอย่างเช่น การ clone
NVIDIA cutlass
repository ลงใน context directory เป็นวิธีที่ดีที่จะให้ agent ค้นหา abstraction ที่เกี่ยวข้องกับ
Layout Algebra, Copy/GEMM atoms, memory hierarchy, example kernels
และอื่นๆ ขณะเขียน kernels ใน CuTeDSL
จากประสบการณ์ของผม เวอร์ชันแรกที่ดีของ kernel จะ compile ผ่านโดยไม่มี error ชัดเจน และผ่านการทดสอบความถูกต้องซึ่งผมจะพูดถึงด้านล่าง
ทดสอบ, Benchmark, และ Profile
เมื่อให้บริบทกับ agent อย่างพอเพียงแล้ว คอขวดที่แท้จริงจะขยับไปที่การตรวจสอบความถูกต้อง (validation) ตัว reference implementation เองและการตรวจสอบเทียบกับมันนั้นสำคัญขึ้นเรื่อยๆ ผมเรียกเฟสนี้ว่า การทดสอบความถูกต้อง หรือเรียกสั้นๆ ว่า การทดสอบ ความเร็วของ reference implementation ไม่สำคัญเท่ากับเจตนาของมัน สิ่งที่คุณตั้งใจจะวัดและตรวจสอบ คือสิ่งที่ agent ของคุณจะมุ่งปรับให้เหมาะสม
โดยปกติ เมื่อการคำนวณไม่ได้ตั้งใจให้เกิดขึ้นในความแม่นยำต่ำ (lower precision) ผมจะวัด Max Absolute/Relative Error (MAE), Mean Squared Error (MSE/RMSE) และ PSNR (Peak Signal to Noise Ratio) แต่เมื่อมีความแม่นยำต่ำเข้ามาเกี่ยวข้อง ผมมักจะวัด PSNR และ Cosine similarity (cossim)
วิธีที่ทำให้ kernel เวอร์ชันต่างๆ รันบน GPU จริงๆ ขึ้นอยู่กับว่า GPU พร้อมใช้งานในเครื่องหรือผ่านคลาวด์ อย่างไรก็ตาม agent ของเราควรมีความสามารถในการเข้าถึง output ของมันไม่ทางใดก็ทางหนึ่ง
ผมพบว่า rung methodology ด้านล่างนี้เป็นวิธีที่ดีในการสร้างฟังก์ชันทดสอบจำนวน N ฟังก์ชัน:
1def rung(name):2 def deco(fn):3 try:4 out = fn()5 results[name] = {"ok": True, **(out or {})}6 print(f"[{name}] ok " + " ".join(f"{k}={v}" for k, v in (out or {}).items()))7 except Exception as e:8 results[name] = {"ok": False, "err": f"{type(e).__name__}: {e}"}9 print(f"[{name}] FAILED {type(e).__name__}: {e}")10 traceback.print_exc()11 return fn12 return deco
ซึ่งคุณสามารถเรียกใช้แบบนี้:
1out = {}23@rung("pre-checks")4def _():5 run_pure_checks()6 run_dsl_checks()78@rung("run")9def _():10 out["o"] = custom_kernel(*inputs)11 torch.cuda.synchronize()12 return {"shape": tuple(out["o"].shape),13 "finite": bool(torch.isfinite(out["o"]).all())}
สำหรับ benchmark rung คุณสามารถทำได้หลายอย่าง:
- Benchmark เวลาการทำงานของ kernel แบบ end-to-end สำหรับเวลาทั้งหมดที่ใช้ไป
- ใช้ intra-kernel tracing เพื่อ benchmark ส่วนต่างๆ ภายใน kernel แล้ว dump ออกมาใน output (โดยใช้ custom tracer หรือ CUPTI)
- Dump ไฟล์ IR, PTX, SASS และ CUBIN ที่สร้างขึ้นไปยัง dumps directory แล้วให้ agent ไล่อ่านมัน
ประเด็นสุดท้ายขยายความได้อีกหน่อย ในบางครั้ง DSL อาจ lower ออกมาเป็น PTX (และสุดท้ายเป็น SASS) ที่ไม่เหมาะสมที่สุด และคุณอาจพบ instruction หรือ shape ที่ดีกว่ามาใช้แทนได้ agent ของเราสามารถอ่านไฟล์ข้อความ PTX/SASS และ inline โค้ดระดับล่างลงไปแทนที่จะปล่อยให้ DSL จัดการส่วนที่ไม่เหมาะสมนั้น อีกครั้ง การส่ง documentation ของ PTX ที่ "ค้นหาได้" เป็นบริบทช่วยได้มากตรงนี้
สิ่งสุดท้ายที่จะรวมทุกอย่างเข้าด้วยกันคือ Profiling ถ้า agent ของคุณเข้าถึง CLI ของ NCU (Nsight Compute Systems) ได้ คุณก็ให้มันทำ profiling และสร้างรายงานสำหรับ kernel ของคุณ ซึ่งเป็นส่วนหนึ่งของ verifiable feedback loop ที่กล่าวถึงข้างต้น
ความคิดปิดท้าย
ถ้าอย่างนั้น การพัฒนา GPU kernel ตายแล้วหรือ?
ก็ใช่ แต่จริงๆ แล้วไม่
ใช่ เพราะส่วนที่ยากอย่าง layouts, indexing, abstractions และโครงสร้างโดยรวมสามารถแก้ได้เป็นส่วนใหญ่โดย agent ที่มีบริบทพอเพียง คุณสามารถลดงาน 2-3 สัปดาห์ให้เหลือ 1-2 วันได้ง่ายๆ และไม่ เพราะคอขวดที่แท้จริงได้ย้ายจาก kernels ไปที่การตรวจสอบความถูกต้องแล้ว ตอนนี้ไม่มีวิธีใดวิธีเดียวที่ถูกต้อง: ยิ่ง context และ harness ของคุณดีเท่าไหร่ กระบวนการก็ยิ่งเร็วขึ้นเท่านั้น เคสเฉพาะทางจะได้ประโยชน์ยิ่งกว่า และที่คุณต้องมีก็แค่ harness ที่ดี
สุดท้ายนี้ แทนที่จะมอง agents เป็นระบบอัตโนมัติ การมองพวกมันเป็นผู้ช่วยที่ฉลาดมากๆ ซึ่งคุณสามารถชี้แนะได้ยังเป็นสิ่งที่จำเป็น นี่คือจุดที่ความเข้าใจพื้นฐานของคุณเกี่ยวกับ GPU และ kernels มีประโยชน์ ส่วนที่เป็นมนุษย์ (คุณ) ยังคงจำเป็นอยู่ตรงนี้
มันเป็นความรู้สึกขมๆ หวานๆ นะ ผมรู้ :)
