ในสมรภูมิการอนุมานโมเดลขนาดใหญ่ ประสิทธิภาพของโอเปอเรเตอร์คือรากฐานของทุกสิ่ง เมื่อพารามิเตอร์ของโมเดล Transformer ทะลุระดับแสนล้าน และสถาปัตยกรรม MoE กลายเป็นมาตรฐานของอุตสาหกรรม การใช้งาน Kernel เพียงรูปแบบเดียวไม่สามารถตอบสนองความต้องการที่หลากหลายของรูปทรง ความแม่นยำ และฮาร์ดแวร์ต่างรุ่นได้อีกต่อไป
- AITER (AI Tensor Engine for ROCm) คือไลบรารีโอเปอเรเตอร์ AI ประสิทธิภาพสูงที่ AMD เปิดตัว ออกแบบมาเพื่อจัดหาโปรแกรมประมวลผลหลัก GPU ที่ผ่านการปรับแต่งอย่างล้ำลึกสำหรับโหลดการอนุมานและการฝึกอบรมบนพื้นฐาน ROCm ไลบรารีนี้รวบรวมโอเปอเรเตอร์ที่พร้อมใช้งานในสภาพแวดล้อมการผลิตหลากหลายชนิด นักพัฒนาเฟรมเวิร์กสามารถผสานรวมเข้ากับสถาปัตยกรรมของตนเองได้อย่างราบรื่น
- https://github.com/ROCm/aiter
- https://rocm.github.io/aiter/
AITER (AI Tensor Engine for ROCm) ที่ AMD เปิดตัวนั้นออกแบบมาเพื่อแก้ปัญหาจุดนี้โดยเฉพาะ มันไม่ใช่แค่ชุดโอเปอเรเตอร์ธรรมดา แต่เป็นเฟรมเวิร์กการจัดตารางเวลาแบบรวมที่ผสานแบ็กเอนด์สี่ประเภท ได้แก่ ASM (Assembly Hand-written Kernel), CK (Composable Kernel) Code Generation, Triton JIT และ FlyDSL ที่พัฒนาขึ้นเอง
AITER ในฐานะที่เป็นฐานโค้ดส่วนกลางของ AMD มอบโอเปอเรเตอร์ประสิทธิภาพสูงต่างๆ สำหรับการเร่งความเร็วโหลดงาน AI และยังเป็นแพลตฟอร์มคุณภาพสูงที่รับข้อกำหนดโอเปอเรเตอร์แบบกำหนดเองทั้งหมดอย่างเป็นหนึ่งเดียว สามารถตอบสนองความต้องการปรับแต่งที่หลากหลาย นักพัฒนาสามารถมุ่งเน้นไปที่การพัฒนาโอเปอเรเตอร์ ในขณะที่ลูกค้าสามารถผสานรวมชุดโอเปอเรเตอร์นี้เข้ากับเฟรมเวิร์กส่วนตัว สาธารณะ หรือเฟรมเวิร์กอื่นๆ ของตนเองได้
ที่สำคัญยิ่งกว่านั้น ภายในของมันได้สร้างวงจรปิดที่สมบูรณ์ตั้งแต่ “การค้นพบรูปทรงที่ยังไม่ได้ปรับแต่ง” ไปจนถึง “การปรับแต่งอัตโนมัติแบบออนไลน์” และ “การโหลดร้อนของ CSV Configuration” เพื่อให้แน่ใจว่าทุกครั้งที่เรียกใช้ GEMM และ MoE จะสามารถเข้าถึง Kernel ที่มีประสิทธิภาพดีที่สุดในขณะรันไทม์
บทความนี้จะพาคุณดำดิ่งสู่แกนกลางของซอร์สโค้ด AITER จากสี่มิติ ได้แก่ การออกแบบสถาปัตยกรรม โครงสร้างพื้นฐานการคอมไพล์แบบ JIT, สายการผลิตหลัก Fused MoE และการปรับแต่งอัตโนมัติของ GEMM
รายละเอียดการผสานรวมของเฟรมเวิร์กนี้แสดงดังตารางด้านล่าง:
| เฟรมเวิร์ก | วิธีการผสานรวม | สถานะ | โอเปอเรเตอร์ที่ใช้ |
|---|---|---|---|
| vLLM | แบ็กเอนด์ Attention เริ่มต้นบนแพลตฟอร์ม ROCm | พร้อมใช้งานในการผลิต | Multi-Head Attention, Multi-Head Linear Attention, Paged Attention, Fused Mixture of Experts, General Matrix Multiply, Root Mean Square Normalization, Rotary Position Embedding + KV Cache |
| SGLang | ผสานรวมเป็นค่าเริ่มต้นในคอนเทนเนอร์อิมเมจ ROCm | พร้อมใช้งานในการผลิต | Attention Mechanism, Fused Mixture of Experts, Block-Scaled Matrix Multiply, All-Reduce Communication, Root Mean Square Normalization |
| ATOM | สร้างขึ้นบนพื้นฐาน AITER โดยตรง | อยู่ระหว่างการพัฒนาอย่างต่อเนื่อง | โอเปอเรเตอร์ AITER ทั้งหมด (Attention, Mixture of Experts, Sampling, Communication) |
| JAX | เชื่อมต่อผ่าน XLA Foreign Function Interface โดยไม่พึ่งพา PyTorch | ระยะทดลอง | Multi-Head Attention / Fused Multi-Head Attention, Root Mean Square Normalization, BF16 Precision Matrix Multiply |
| เอ็นจิ้นการอนุมานที่พัฒนาขึ้นเองของลูกค้าหลายราย | ผสานรวมระดับ Kernel อย่างล้ำลึก | พร้อมใช้งานในการผลิต | Attention, Mixture of Experts, General Matrix Multiply, Quantization Operators |
ประสิทธิภาพของโอเปอเรเตอร์หลักแสดงดังนี้ สามารถดูรายละเอียดเพิ่มเติมได้ที่ ATOM/benchmark-dashboard[1]:
| โอเปอเรเตอร์ | อัตราเร่งความเร็ว |
|---|---|
| Multi-Head Linear Attention Decode Kernel (MLA decode kernel) | สูงสุด 17 เท่า |
| Multi-Head Attention Prefill Kernel (MHA prefill kernel) | สูงสุด 14 เท่า |
| Block-Scaled Fused Mixture of Experts | สูงสุด 3 เท่า |
| Block-Scaled General Matrix Multiply | สูงสุด 2 เท่า |
| DeepSeek-R1 End-to-End Inference (SGLang) | 6,484 → 13,744 โทเค็น/วินาที (2.1 เท่า) |
| JAX-AITER Attention Operator (แพลตฟอร์ม MI350) | ความเร็วที่เพิ่มขึ้นมัธยฐาน 4.39 เท่า |
รุ่น GPU ของ AMD ที่รองรับในปัจจุบันมีดังนี้:
| รุ่น GPU | สถาปัตยกรรม | สถานะการรองรับ |
|---|---|---|
| AMD Instinct MI300X | gfx942 (CDNA3) | รองรับอย่างสมบูรณ์ |
| AMD Instinct MI325X | gfx942 (CDNA3) | รองรับอย่างสมบูรณ์ |
| AMD Instinct MI350 | gfx950 (CDNA4) | รองรับ |
| AMD Instinct MI355X | gfx950 (CDNA4) | รองรับ |
unsetunsetสารบัญunsetunset
- เริ่มต้นใช้งานอย่างรวดเร็ว
- หนึ่ง ภาพรวมสถาปัตยกรรมและปรัชญาการออกแบบ
- 1.1 ทางเข้าเดียว สี่แบ็กเอนด์
- 1.2 กระแสข้อมูลหลัก
- สอง โครงสร้างพื้นฐานการคอมไพล์ JIT: สายการผลิตอัตโนมัติจากซอร์สโค้ดสู่ .so
- 2.1 ตัวตกแต่ง compile_ops: ประกาศก็คือคอมไพล์
- 2.2 build_module: เอ็นจิ้นการคอมไพล์ที่ปลอดภัยต่อหลายกระบวนการ
- 2.3 torch_compile_guard: เชื่อมต่อกับ torch.compile ได้อย่างราบรื่น
- สาม Fused MoE: การจัดเรียง Token และการประสานงาน GEMM หลายขั้นตอนอย่างแม่นยำ
- 3.1 การจัดเรียง Token: ปูทางให้ GPU ทำงานแบบขนาน
- 3.2 หนึ่งขั้นตอน vs สองขั้นตอน: เอ็นจิ้นการตัดสินใจแบบไดนามิก
- 3.3 การปรับ Block Size แบบปรับตัวได้: ใช้ทุก CU ให้คุ้มค่า
- 3.4 K-Split: ตัวช่วยอัตราการใช้ CU ในสถานการณ์จำนวน Token น้อย
- สี่ การปรับแต่ง GEMM อัตโนมัติ: วงจรปิดจากการค้นพบรูปทรงสู่ Kernel ที่ดีที่สุด
- 4.1 ตารางเส้นทางหลายแบ็กเอนด์
- 4.2 Skinny GEMM: อาวุธลับในสถานการณ์ M เล็ก
- 4.3 กลไกสองรางของการปรับแต่งออนไลน์และออฟไลน์
- ห้า ระบบนิเวศ Quantization: ครอบคลุมทุกความแม่นยำตั้งแต่ FP4 ถึง INT8
- 5.1 การแบ่งระดับประเภท Quantization อย่างละเอียด
- 5.2 การจัดเรียงแบบ Fused ของ MXFP4 Quantization
- หก จุดที่น่าสนใจ
- 6.1 การใช้ LRU Cache อย่างแพร่หลาย
- 6.2 กลไกการถอยกลับของ FlyDSL
- 6.3 ctypes FFI: เส้นทางการเรียกที่ไม่ผ่าน Torch โดยเลี่ยง PyTorch
- บทสรุป
unsetunsetเริ่มต้นใช้งานอย่างรวดเร็วunsetunset
ขั้นตอนการติดตั้ง AITER นั้นง่ายมาก เพียงรันคำสั่งสามคำสั่งต่อไปนี้ก็เสร็จสมบูรณ์:
git clone –recursive https://github.com/ROCm/aiter.git
cd aiter
python3 setup.py develop
หากต้องการเปิดใช้งานการรองรับ FlyDSL Mixed Precision MoE จำเป็นต้องรันคำสั่งติดตั้งเพิ่มเติมอีกหนึ่งคำสั่ง:
pip install –pre flydsl
สำหรับข้อมูลเพิ่มเติมเกี่ยวกับตัวเลือกการติดตั้งอื่นๆ (เช่น Triton Communication Library, Iris Library ฯลฯ) โปรดดูไฟล์ README.md ในไดเรกทอรีรากของโปรเจกต์ หลังจากติดตั้งเสร็จ นักพัฒนาสามารถเรียกใช้โอเปอเรเตอร์ทั้งหมดในสภาพแวดล้อม Python ผ่านคำสั่ง import aiter ตัวอย่างเช่น รันสคริปต์ทดสอบในตัวเพื่อตรวจสอบว่าการติดตั้งสำเร็จหรือไม่:
python3 op_tests/test_layernorm2d.py
unsetunsetหนึ่ง ภาพรวมสถาปัตยกรรมและปรัชญาการออกแบบunsetunset
1.1 ทางเข้าเดียว สี่แบ็กเอนด์
แนวคิดการออกแบบหลักที่สุดของ AITER สามารถสรุปได้เป็นประโยคเดียว: ใช้ API แบบรวมในเลเยอร์ Python เพื่อปิดบังความแตกต่างของ Kernel ภายใน และเลือกใช้แบ็กเอนด์ที่เหมาะสมที่สุดในขณะรันไทม์แบบไดนามิกผ่านการขับเคลื่อนด้วย Configuration
จากโครงสร้างการนำเข้าของไฟล์ aiter/__init__.py เราสามารถมองเห็นวิธีการใช้งานการออกแบบแบบ分层 (layered) นี้ได้อย่างชัดเจน:
ที่มา: aiter/init.py
from .ops.gemm_op_a8w8 import * # INT8/FP8 Quantized GEMM
from .ops.gemm_op_a16w16 import * # BF16/FP16 GEMM
from .ops.gemm_op_a4w4 import * # FP4 Quantized GEMM
from .ops.moe_op import * # MoE Operator
from .ops.attention import * # MHA/MLA/PA
from .ops.activation import * # Activation Functions เช่น SiLU/GeLU
from .ops.rope import * # RoPE Positional Encoding
แต่ละโมดูล ops จะเปิดเผยลายเซ็นฟังก์ชันแบบรวม ในขณะที่ภายใน ระบบจะทำการกำหนดเส้นทางคำขอไปยังการใช้งานแบ็กเอนด์ที่แตกต่างกันสี่แบบแบบไดนามิก ตามหลายมิติ เช่น สถาปัตยกรรมฮาร์ดแวร์ (เช่น gfx942/gfx950), รูปทรงเมทริกซ์ (M/N/K), ประเภท Quantization (เช่น per_Token/per_1x128/per_1x32) และการกำหนดค่าที่ปรับแต่งไว้ล่วงหน้า (เก็บในไฟล์ CSV) แบ็กเอนด์ทั้งสี่นี้ได้แก่:
- ASM Assembly Kernel: ใช้ GCN/RDNA Assembly ที่เขียนด้วยมือ ประสิทธิภาพสูงสุดแต่จำกัดเฉพาะรูปทรงที่กำหนด
- CK (Composable Kernel): 基于 AMD โอเพนซอร์สเฟรมเวิร์กการสร้างโค้ด C++ แบบเทมเพลต
- Triton: Kernel Python DSL ที่พัฒนาบนพื้นฐาน OpenAI Triton
- FlyDSL: ภาษาเฉพาะโดเมน (DSL) ที่ AMD พัฒนาขึ้นเอง ออกแบบมาสำหรับสถานการณ์ MoE แบบ Mixed Precision โดยเฉพาะ
1.2 กระแสข้อมูลหลัก
กระแสข้อมูลของการเรียกใช้ Fused MoE Inference ทั่วไปมีดังนี้:
ผู้ใช้เรียก fused_moe()
→ การ推断ประเภท Quantization และการจัดแนวรูปทรง
→ get_2stage_cfgs() ค้นหา CSV Configuration
→ moe_sorting() จัดเรียง Token ตาม Expert
→ ดำเนินการ GEMM 1-stage หรือ 2-stage
→ ส่งออก moe_buf
ทุกขั้นตอนในสายการผลิตนี้ประกอบด้วยการตัดสินใจทางวิศวกรรมที่精巧 ต่อไปเราจะวิเคราะห์เชิงลึกทีละขั้นตอน
unsetunsetสอง โครงสร้างพื้นฐานการคอมไพล์ JIT: สายการผลิตอัตโนมัติจากซอร์สโค้ดสู่ .sounsetunset
2.1 ตัวตกแต่ง compile_ops: ประกาศก็คือคอมไพล์
ระบบ JIT ของ AITER เป็นโครงสร้างหลักของทั้งโปรเจกต์
นักพัฒนา只需ใช้ @compile_ops ตกแต่งฟังก์ชันว่างที่มีเพียงลายเซ็นประเภท เฟรมเวิร์กจะดำเนินการทั้งกระบวนการโดยอัตโนมัติ: ค้นหาไฟล์ .so ที่คอมไพล์ไว้ล่วงหน้า หากไม่พบจะ触发 JIT Compilation, โหลดโมดูลที่สร้างขึ้น, และสุดท้ายเรียกใช้ C++ Operator
ที่มา: aiter/ops/gemm_op_a8w8.py
@compile_ops(
“module_gemm_a8w8″, fc_name=”gemm_a8w8”,
gen_fake=gen_gemm_a8w8_ck_fake_tensors
)
def gemm_a8w8_ck(
XQ: torch.Tensor, WQ: torch.Tensor,
x_scale: torch.Tensor, w_scale: torch.Tensor,
Out: torch.Tensor, bias: Optional[torch.Tensor] = None,
splitK: int = 0,
) -> torch.Tensor: …
โปรดสังเกตว่า เนื้อหาของฟังก์ชันคือ ... (Ellipsis) — มันไม่จำเป็นต้องมีการใช้งาน Python เลย ฟังก์ชัน wrapper ภายใน compile_ops จะดำเนินการตามขั้นตอนต่อไปนี้:
- พยายามโหลดโมดูลที่คอมไพล์แล้วผ่าน
get_module(md_name) - หากเกิดข้อยกเว้น
ModuleNotFoundErrorจะเรียกbuild_module()เพื่อเริ่มกระบวนการคอมไพล์ - หลังจากคอมไพล์เสร็จ จะรับพอยน์เตอร์ของฟังก์ชัน C++ ผ่าน
getattr(module, loadName)
2.2 build_module: เอ็นจิ้นการคอมไพล์ที่ปลอดภัยต่อหลายกระบวนการ
build_moduleคือแกนหลักหนักของระบบ JIT ซึ่งต้องแก้ปัญหาสำคัญหลายประการ
ล็อคการคอมไพล์หลายกระบวนการ: ในสถานการณ์การฝึกอบรมหลาย GPU หลายกระบวนการอาจ触发การคอมไพล์โมดูลเดียวกันพร้อมกัน AITER ใช้กลไก File Lock (FileBaton) เพื่อให้แน่ใจว่ามีเพียงกระบวนการเดียวเท่านั้นที่ดำเนินการคอมไพล์:
ที่มา: aiter/jit/core.py
def mp_lock(lockPath, MainFunc, FinalFunc=None, WaitFunc=None):
baton = FileBaton(lockPath)
if baton.try_acquire():
try:
ret = MainFunc()
finally:
if FinalFunc is not None:
FinalFunc()
baton.release()
else:
baton.wait() # กระบวนการอื่นรอ
โอเค โปรดตรวจสอบส่วนของบทความที่ถูกเขียนใหม่เชิงลึกตามความต้องการของคุณ
การปรับตัวของ Compile Flags
เฟรมเวิร์กนี้มีความสามารถในการตรวจจับเวอร์ชันของ HIP Compiler โดยอัตโนมัติ และสามารถเพิ่ม Optimization Flags ได้ทีละรายการตามนั้น โดยเฉพาะอย่างยิ่ง เมื่อรันบน ROCm 6.2 ขึ้นไป มันจะเปิดใช้งาน Flags amdgpu-early-inline-all และ amdgpu-function-calls=false นอกจากนี้ สำหรับสถาปัตยกรรม gfx950 เฟรมเวิร์กจะเพิ่ม Flag ที่รองรับ数据类型 FP4 โดยอัตโนมัติ ซึ่งมีตรรกะการทำงานดังนี้:
# ที่มา: aiter/jit/core.py (ภายใน build_module)
if hip_version > Version("6.2.41132"):
flags_hip += [
"-mllvm -amdgpu-early-inline-all=true",
"-mllvm -amdgpu-function-calls=false",
]
if get_gfx() == "gfx950" and int(os.getenv("AITER_FP4x2", "1")) > 0:
flags_hip += ["-D__Float4_e2m1fn_x2"]
2.3 torch_compile_guard: ผสานรวมกับ torch.compile ได้อย่างราบรื่น
ใน AITER ทุกโอเปอเรเตอร์จะถูกลงทะเบียนเป็น Custom Operator ภายใต้ torch.ops.aiter.* ผ่านตัวตกแต่ง torch_compile_guard กลไกนี้ทำให้แน่ใจว่าพวกมันสามารถถูก捕获และปรับให้เหมาะสมโดยกราฟการคำนวณของ torch.compile ตัวตกแต่งนี้สามารถ推导 Schema ของโอเปอเรเตอร์ สร้าง FakeTensor สำหรับการติดตามเชิงสัญลักษณ์ (symbolic tracing) และจัดการตรรกะการกระจายทั้งฝั่ง CUDA และ CPU ได้อย่างเหมาะสม:
# ที่มา: aiter/jit/utils/torch_guard.py
aiter_lib.define(op_schema, tags=tags)
aiter_lib.impl(f"aiter::{loadName}", custom_func, dispatch_key="CUDA")
aiter_lib.impl(f"aiter::{loadName}", custom_func, dispatch_key="CPU")
aiter_lib._register_fake(f"{loadName}", fake_func)
ด้วยเหตุนี้ ผู้ใช้สามารถใช้โอเปอเรเตอร์ AITER ในบริบทของ torch.compile(model) ได้โดยตรง คอมไพเลอร์จะสามารถจัดการรายละเอียดระดับล่าง เช่น การ推导รูปทรงและการจัดสรรหน่วยความจำได้อย่างถูกต้อง
สาม Fused MoE: การจัดเรียง Token และการประสานงาน GEMM หลายขั้นตอนอย่างแม่นยำ
Fused MoE เป็นโอเปอเรเตอร์ที่มีความซับซ้อนสูงที่สุดและสำคัญที่สุดใน AITER ไฟล์การใช้งานหลัก fused_moe.py มีจำนวนโค้ดเกือบ 2,000 บรรทัด เป้าหมายหลักของมันคือการรวมสายการผลิต MoE ที่สมบูรณ์ “การเลือกเส้นทาง Gating → การจัดเรียง Token ใหม่ → การคำนวณ GEMM ของ Expert → การรวมแบบถ่วงน้ำหนัก” ให้เป็นการเปิดใช้งาน Kernel น้อยที่สุดเท่าที่จะเป็นไปได้
3.1 การจัดเรียง Token: ปูทางให้ GPU ทำงานแบบขนาน
ขั้นตอนแรกของการอนุมาน MoE คือการจัดเรียง Token ใหม่ ตาม ID ของ Expert ที่它们ถูกจัดสรรให้ การทำเช่นนี้เพื่อให้แน่ใจว่า Token ที่ถูกประมวลผลโดย Expert เดียวกันนั้นอยู่ในหน่วยความจำที่ต่อเนื่องกัน ทำให้การดำเนินการ GEMM ที่ตามมาสามารถทำงานในรูปแบบการคูณเมทริกซ์แบบ密集 (dense) ที่มีประสิทธิภาพสูง
# ที่มา: aiter/fused_moe.py
def _moe_sorting_impl(...):
M, topk = topk_ids.shape
max_num_tokens_padded = topk_ids.numel() + num_experts * block_size - topk
max_num_m_blocks = (max_num_tokens_padded + block_size - 1) // block_size
sorted_ids = torch.empty(max_num_tokens_padded, dtype=dtypes.i32, device=device)
sorted_weights = torch.empty(max_num_tokens_padded, dtype=dtypes.fp32, device=device)
sorted_expert_ids = torch.empty(max_num_m_blocks, dtype=dtypes.i32, device=device)
fwd_fn = aiter.moe_sorting_opus_fwd if use_opus else aiter.moe_sorting_fwd
fwd_fn(topk_ids, topk_weights, sorted_ids, sorted_weights, ...)
ในที่นี้ใช้ กลยุทธ์ Padding ที่ชาญฉลาด: max_num_tokens_padded ไม่เพียงรวมจำนวน Token ทั้งหมดคูณด้วย TopK ที่展开แล้ว แต่ยัง预留พื้นที่จัดเรียงขนาด block_size เพิ่มเติมสำหรับแต่ละ Expert ข้อดีของการออกแบบนี้คือ สามารถ确保การดำเนินการ GEMM แบบ Tile-based ในภายหลังเมื่อจัดการกับ Token ท้ายๆ จะไม่ต้อง编写ตรรกะการจัดการขอบเขตพิเศษเนื่องจากจำนวน Token ไม่เพียงพอต่อหนึ่ง Block
3.2 หนึ่งขั้นตอน vs สองขั้นตอน: เอ็นจิ้นการตัดสินใจแบบไดนามิก
AITER มีเส้นทางการดำเนินการสองแบบสำหรับ MoE:
- 1-stage: ดำเนินการ Gate+Up Projection, การคำนวณฟังก์ชัน激活 และ Down Projection ภายใน Kernel เดียว
- 2-stage: แยกการคำนวณเป็นสอง Kernel — ขั้นตอนแรกดำเนินการ Gate+Up+Activation ขั้นตอนที่สองดำเนินการ Down Projection
การเลือกเส้นทางใดนั้นถูกกำหนดโดยฟังก์ชัน get_2stage_cfgs() ซึ่งถือได้ว่าเป็น “สมอง” ของตรรกะการจัดตารางเวลา MoE ทั้งหมด เกณฑ์การตัดสินใจหลักประกอบด้วยประเด็นต่อไปนี้:
# ที่มา: aiter/fused_moe.py (ภายใน get_2stage_cfgs)
# เงื่อนไขการ启用เส้นทางหนึ่งขั้นตอนแตกต่างกันไปตามประเภท Quantization
if q_type == QuantType.per_1x128:
run_1stage = token > 32 and (inter_dim % 128 == 0)
elif q_type == QuantType.per_Token and q_dtype_w == dtypes.fp8:
run_1stage = token > 16 or inter_dim % 128 != 0
elif q_type != QuantType.per_1x32:
run_1stage = token < 256
สัญชาตญาณการออกแบบเบื้องหลังนี้คือ:
- เมื่อจำนวน Token น้อย (เช่น สถานการณ์ decode) การใช้โหมด 1-stage สามารถลดค่าใช้จ่ายในการเปิดใช้งาน Kernel ได้อย่างมีประสิทธิภาพ
- เมื่อจำนวน Token มาก (เช่น สถานการณ์ prefill) โหมด 2-stage อนุญาตให้ผลลัพธ์ระหว่างกลางได้รับการ Quantization ที่ละเอียดยิ่งขึ้น (เช่น MXFP4 Dynamic Quantization) ระหว่างการดำเนินการ GEMM สองครั้ง ทำให้ได้ความหนาแน่นในการคำนวณที่สูงขึ้น
3.3 การปรับ Block Size แบบปรับตัวได้: ใช้ทุกหน่วยคำนวณให้คุ้มค่า
ฟังก์ชัน get_block_size_M แสดงให้เห็นถึงกลยุทธ์การจัดตารางเวลาที่รับรู้ถึงฮาร์ดแวร์อย่าง极致:
# ที่มา: aiter/fused_moe.py
@functools.lru_cache(maxsize=2048)
def get_block_size_M(token, topk, expert, inter_dim):
cu_num = get_cu_num() # จำนวน Compute Unit ของ GPU
tileN = 128
tgN = (inter_dim + tileN - 1) // tileN
support_list = [32, 64, 128]
tmp = []
for el in support_list:
max_num_tokens = token * topk + expert * el - topk
tg_num = tgN * (max_num_tokens + el - 1) // el # จำนวน Tile Group ทั้งหมด
rnd = (tg_num + cu_num - 1) // cu_num # จำนวนรอบที่แต่ละ CU ต้องดำเนินการ
empty = cu_num - tg_num % cu_num # จำนวน CU ที่ว่าง
tmp.append((rnd, empty, el))
return sorted(tmp, key=lambda x: x[:2])[0][-1]
ฟังก์ชันนี้จะ遍历 Block Size ที่เป็นไปได้ทั้งหมด (32, 64, 128) สำหรับแต่ละ Configuration จะคำนวณ จำนวน Tile ทั้งหมด → อัตราการใช้ CU → จำนวน CU ที่ว่าง และ最终เลือกแผนที่ดีที่สุดที่มี “จำนวนรอบการดำเนินการน้อยที่สุด และ CU ว่างน้อยที่สุด” ซึ่งเทียบเท่ากับการวิเคราะห์ GPU Occupancy แบบ精简ในเลเยอร์ Python
3.4 K-Split: ตัวช่วยอัตราการใช้ CU ในสถานการณ์จำนวน Token น้อย
ในสถานการณ์ decode (จำนวน Token น้อยมาก เช่น 1-8) มิติ M ของ GEMM MoE มีขนาดเล็กมาก ส่งผลให้จำนวน Tile ทั้งหมดไม่เพียงพอที่จะใช้ CU ทั้งหมดให้เกิดประโยชน์ ฟังก์ชัน get_ksplit สร้างความขนานเพิ่มเติมโดยการแบ่งการคำนวณในมิติ K:
# ที่มา: aiter/fused_moe.py
def get_ksplit(token, topk, expert, inter_dim, model_dim):
if token * topk > expert:
return 0 # Token เพียงพอ ไม่จำเป็นต้อง split
cu_num = get_cu_num()
tg_num = tgN * tgM
if tg_num >= cu_num:
return 0 # CU เต็มแล้ว
split_max = (cu_num + tg_num - 1) // tg_num
for i in reversed(range(2, split_max + 1)):
if (model_dim % i == 0) and ((model_dim // i) % 256 == 0):
return i
return 0
เปรียบเสมือนการตัดเค้กชิ้นใหญ่ (การดำเนินการสะสมในมิติ K) ออกเป็นหลายชิ้น ให้คนมากขึ้น (CU) จัดการพร้อมกัน แล้วค่อยรวมผลลัพธ์ — นี่คือกลยุทธ์ Split-K แบบคลาสสิก
สี่ การปรับแต่ง GEMM อัตโนมัติ: วงจรปิดจากการค้นพบรูปทรงสู่ Kernel ที่ดีที่สุด
4.1 ตารางเส้นทางหลายแบ็กเอนด์
GEMM BF16/FP16 ของ AITER รองรับแบ็กเอนด์ห้าประเภท สลับผ่านตารางเส้นทางที่简洁:
# ที่มา: aiter/tuned_gemm.py
solMap = {
"torch": torch_gemm, # PyTorch F.linear สำรอง
"hipblaslt": hipb_gemm, # ไลบรารี hipBLASLt
"skinny": skinny_gemm, # Assembly เฉพาะเมทริกซ์แคบพิเศษ
"asm": asm_gemm, # Assembly Kernel ที่เขียนด้วยมือ
"triton": triton_gemm, # Triton JIT Kernel
}
การตัดสินใจเส้นทางถูกขับเคลื่อนโดยไฟล์ Configuration CSV อย่างสมบูรณ์ ฟังก์ชัน get_GEMM_A16W16_config จะใช้ (cu_num, M, N, K, dtype, otype, ...) เป็นดัชนีในการค้นหาตาราง:
# ที่มา: aiter/tuned_gemm.py
config = cfg.get(
(cu_num, padded_M, N, K, bias, str(dtype), str(otype), scaleAB, bpreshuffle),
None,
)
หากรูปทรงที่แน่นอนไม่ตรง ระบบจะลองใช้กลยุทธ์ Padding สองแบบ (gl=0 และ gl=1 ของ get_padded_m) เพื่อการจับคู่แบบ模糊
4.2 Skinny GEMM: อาวุธลับในสถานการณ์ M เล็ก
เมื่อ M มีขนาดเล็กมาก (1-16) และ N ไม่เกิน 1-2 เท่าของจำนวน CU AITER จะข้ามไลบรารี GEMM ทั่วไป และใช้ Assembly Kernel Split-K ที่เขียนด้วยมือโดยตรง:
# ที่มา: aiter/tuned_gemm.py (ตรรกะเส้นทางเริ่มต้น)
if (
((M == 1 and N <= 2 * cu_num) or (M > 1 and M <= 4 and N <= cu_num))
and K <= 9216
or (M > 4 and M <= 8 and N <= cu_num) and K <= 5120
):
default_config["libtype"] = "skinny"
default_config["solidx"] = 2
Kernel เหล่านี้ได้รับการปรับให้เหมาะสมสำหรับ 304 CU ของ MI300X โดยเฉพาะ ในสถานการณ์ decode (M=1) มักจะเร็วกว่าไลบรารีทั่วไป 2-3 เท่า
4.3 กลไกสองรางของการปรับแต่งออนไลน์และออฟไลน์
เลเยอร์ MoE ของ AITER รองรับการปรับแต่งออนไลน์: เมื่อรันไทม์พบชุดรูปทรงที่ยังไม่ได้บันทึกใน CSV และตั้งค่าตัวแปรสภาพแวดล้อม
AITER_ONLINE_TUNE=1เฟรมเวิร์กจะ触发กระบวนการปรับแต่งโดยอัตโนมัติ:
# ที่มา: aiter/fused_moe
⚠️ หมายเหตุ: เนื้อหาได้รับการแปลโดย AI และตรวจสอบโดยมนุษย์ หากมีข้อผิดพลาดโปรดแจ้ง
☕ สนับสนุนค่ากาแฟทีมงาน
หากคุณชอบบทความนี้ สามารถสนับสนุนเราได้ผ่าน PromptPay
SCAN TO PAY WITH ANY BANK
本文来自网络搜集，不代表คลื่นสร้างอนาคต立场，如有侵权，联系删除。转载请注明出处：https://www.itsolotime.com/th/archives/32381