ในสนามประลองการอนุมานของโมเดลภาษาขนาดใหญ่ ประสิทธิภาพการคำนวณของกลไก Attention เป็นตัวกำหนดเพดานปริมาณงานของระบบโดยตรง ในฐานะตัวแปรใหม่ของตระกูล Linear Attention Gated DeltaNet (GDN) ได้ใช้กลยุทธ์การคำนวณแบบแบ่งส่วน (chunk-wise recurrence) อย่างชาญฉลาด เพื่อหลีกเลี่ยงคอขวดของความซับซ้อนระดับกำลังสองของ Standard Attention ตั้งแต่รากฐาน
อย่างไรก็ตาม “ความเป็นเชิงเส้น” ในระดับอัลกอริทึมไม่ได้แปลว่า “ประสิทธิภาพสูง” ในระดับฮาร์ดแวร์โดยอัตโนมัติ เมื่อขั้นตอนการคำนวณที่เชื่อมโยงกันอย่างแน่นหนาทั้งเจ็ดขั้นตอนถูกแยกออกเป็นการเรียกใช้เคอร์เนลอิสระเจ็ดครั้ง ค่าใช้จ่ายในการจัดตารางเวลาของเลเยอร์ Python ความหน่วงในการซิงโครไนซ์ระหว่างคอร์ และการดำเนินการอ่าน/เขียน Global Memory ที่บ่อยครั้ง ก็เพียงพอที่จะทำให้อัตราเร่งทางทฤษฎีหมดไป
- Megakernel สำหรับ Gated DeltaNet, ปรับแต่งโดย PTO-ISA
- https://github.com/huawei-csl/megagdn-pto
- 7000 คำ, อ่าน 38 นาที, พอดแคสต์ 37 นาที
โปรเจกต์ MegaGDN-PTO จากทีม Huawei CSL ใช้ PTO-ISA ซึ่งเป็นชุดนามธรรมของชุดคำสั่งระดับล่างของ Ascend AI Core โดยตรง เขียน Megakernel ที่รวมทั้งเจ็ดขั้นตอนไว้ในการเรียกใช้เคอร์เนลครั้งเดียว บนชิป Atlas 910B โปรเจกต์นี้มีประสิทธิภาพเหนือกว่าเกณฑ์มาตรฐาน Triton-Ascend ถึง 1.5 ถึง 3 เท่า พร้อมรับประกันความแม่นยำของโมเดลที่ไม่สูญเสียเลย
แผนภูมิแท่งนี้แสดงประสิทธิภาพความหน่วงของเคอร์เนลในขั้นตอนเดียวภายใต้การกำหนดค่าเดียวกัน (16 ลำดับ, ความยาว 8192 โทเค็น) ของสามวิธีการนำไปใช้: การนำ PTO ดั้งเดิมใช้เวลาทั้งหมดประมาณ 24ms โดยแต่ละขั้นตอนใช้เวลาใกล้เคียงกัน; การนำ Triton (BT=64) ใช้เวลาสูงสุดถึง 52ms โดยที่ขั้นตอน chunk h และ chunk o ใช้สัดส่วนเวลามากที่สุด; การนำ Triton (BT=128) ใช้เวลาทั้งหมดประมาณ 22ms ประสิทธิภาพโดยรวมดีกว่าเวอร์ชัน BT=64 แต่บางขั้นตอน (เช่น scaled dot KKT, solve tril, chunk o) มีเครื่องหมาย “Triton stage failed” แสดงว่าการใช้งานโอเปอเรเตอร์บางตัวมีปัญหาความเข้ากันได้ โดยรวมแล้ว การนำ PTO ดั้งเดิมมีความเสถียรที่สุด ส่วน Triton ที่ BT=128 มีศักยภาพด้านประสิทธิภาพ แต่โอเปอเรเตอร์มีความเสถียรไม่เพียงพอ ในขณะที่การนำ Triton ที่ BT=64 มีประสิทธิภาพลดลงอย่างเห็นได้ชัด แผนภูมิแท่งสองคอลัมน์นี้เปรียบเทียบความแม่นยำของ PTO megakernel กับ Triton baseline บนโมเดลต่างๆ กราฟซ้ายแสดงค่า Perplexity ของ Wikitext (ยิ่งต่ำยิ่งดี) พบว่าค่าของทั้งสองบนโมเดลซีรีส์ Qwen3.5-0.8B/9B, Qwen3.6-27B/35B เกือบทับกันสนิท โดยค่า Perplexity ลดลงอย่างมีนัยสำคัญเมื่อขนาดโมเดลเพิ่มขึ้น กราฟขวาแสดงความแม่นยำของ MMLU (ยิ่งสูงยิ่งดี) ก็มีความสอดคล้องสูงเช่นกัน โดยค่าความแตกต่างของความแม่นยำระหว่าง PTO megakernel กับ Triton baseline ในแต่ละโมเดลอยู่ภายใน 1% และเพิ่มขึ้นอย่างต่อเนื่องตามขนาดโมเดลที่ใหญ่ขึ้น โดยรวมแล้ว PTO megakernel ในขณะที่ปรับปรุงประสิทธิภาพการคำนวณอย่างมาก ก็สามารถบรรลุความแม่นยำของโมเดลที่เท่ากันทุกประการกับ Triton baseline ซึ่งยืนยันถึงผลการปรับแต่งที่ไม่สูญเสียความแม่นยำ แผนภูมินี้แสดงการเปรียบเทียบประสิทธิภาพระหว่าง PTO megakernel กับ Triton baseline ในช่วง Prefill ของโมเดล Qwen36-35B ในช่วงความยาวพรอมต์ตั้งแต่ 512 ถึง 65536 โทเค็น PTO megakernel บรรลุอัตราเร่ง 1.1-1.25 เท่า โดยจุดสูงสุดอยู่ที่ 4096 โทเค็น; ในส่วนของความหน่วง TTFT PTO ต่ำกว่า Triton ในทุกความยาว โดยเฉพาะที่ 65536 โทเค็น ช่องว่างความหน่วงเพิ่มขึ้นอย่างมีนัยสำคัญ; ในด้านปริมาณงาน ความเร็วในการสร้างโทเค็นของ PTO นำตลอด โดยถึงจุดสูงสุดประมาณ 16000 โทเค็น/วินาทีที่ 16384 โทเค็น โดยข้อได้เปรียบด้านประสิทธิภาพโดยรวมจะชัดเจนยิ่งขึ้นเมื่อความยาวพรอมต์เพิ่มขึ้น โดยรวมแล้ว PTO megakernel เหนือกว่า Triton baseline ในช่วง Prefill อย่างสมบูรณ์ โดยเฉพาะในสถานการณ์ลำดับยาวที่แสดงข้อได้เปรียบด้านประสิทธิภาพที่แข็งแกร่งกว่า ชุดกราฟย่อยสามแถวนี้เปรียบเทียบความแตกต่างด้านประสิทธิภาพระหว่าง PTO megakernel กับ Triton baseline ในช่วง Prefill ของโมเดล Qwen36-35B เส้นโค้งอัตราเร่งในส่วนบนแสดงให้เห็นว่า PTO มีข้อได้เปรียบด้านความเร่งที่เสถียร 1.1-1.18 เท่า; เส้นโค้งความหน่วง TTFT ตรงกลางแสดงให้เห็นว่าความหน่วงการตอบสนองแพ็กเกจแรกของ PTO ต่ำกว่า Triton ตลอด โดยเฉพาะที่ 32768 โทเค็นที่เห็นความแตกต่างชัดเจน; เส้นโค้งปริมาณงานด้านล่างแสดงให้เห็นว่าความเร็วในการสร้างโทเค็นของ PTO นำตลอด โดยถึงจุดสูงสุดประมาณ 19000 โทเค็น/วินาทีที่ 16384 โทเค็น โดยรวมแล้ว PTO megakernel เหนือกว่า Triton baseline ในช่วง Prefill อย่างสมบูรณ์ โดยเฉพาะในสถานการณ์ลำดับขนาดกลางถึงยาวที่มีความหน่วงต่ำกว่าและปริมาณงานสูงกว่า แสดงให้เห็นถึงข้อได้เปรียบด้านประสิทธิภาพที่เสถียร
unsetunsetสารบัญunsetunset
ในฐานะบรรณาธิการบริหารอาวุโสด้านบทความเทคนิคและผู้เชี่ยวชาญด้าน “การเขียนบทความใหม่” ระดับสูง ผมจะดำเนินการเขียนใหม่และลดความซ้ำซ้อนของส่วนบทความที่คุณให้มาอย่างลึกซึ้งทันที
- เริ่มต้นใช้งานอย่างรวดเร็ว
- หนึ่ง ภาพรวมสถาปัตยกรรมและปรัชญาการออกแบบ
- 1.1 สายการผลิตการคำนวณเจ็ดขั้นตอนของ GDN
- 1.2 คอขวดประสิทธิภาพของการดำเนินการแบบแยกขั้นตอน
- 1.3 วิธีการแก้ปัญหาของ Megakernel
- 1.4 การจัดระเบียบโค้ดโปรเจกต์
- สอง โมเดลการเขียนโปรแกรมฮาร์ดแวร์ NPU: Vec, Cube และไปป์ไลน์แบบชัดแจ้ง
- 2.1 สถาปัตยกรรมภายใน AI Core
- 2.2 ลำดับชั้นหน่วยความจำและนามธรรม Tile
- 2.3 ภาพรวมของดั้งเดิมการซิงโครไนซ์
- สาม ระบบคอมไพล์: การปรับแต่งเฉพาะช่วงคอมไพล์ของ Bisheng JIT
- 3.1 กลยุทธ์การแทรกค่าคงที่ในช่วงคอมไพล์
- 3.2 การคอมไพล์แบบเพิ่มหน่วยและแคช
- 3.3 การปรับตัวตามฮาร์ดแวร์
- สี่ การวิเคราะห์เคอร์เนลหลักทีละระดับ
- 4.1 chunk_cumsum: ผลรวมนำหน้าของ Vec-only
- 4.2 scaled_dot_kkt: ความร่วมมือแบบดูอัลคอร์ของ Cube และ Vec
- 4.3 wy_fast: ความร่วมมือ Vec+Cube ของการแยกส่วน WY
- 4.4 solve_tril: การย้อนกลับเมทริกซ์สามเหลี่ยมแบบเรียกซ้ำของ CubeCore
- 4.5 จุดเข้าแบบรวม Megakernel: การบูรณาการเจ็ดขั้นตอน
- 4.6 การจับมือสามขั้นตอนของสิ่งกีดขวางต่างชนิด SyncAllImpl
- ห้า เลเยอร์อินเทอร์เฟซ Python: ห่วงโซ่ที่สมบูรณ์ตั้งแต่การคอมไพล์ไปจนถึงการเรียก
- 5.1 การเรียกแบบคัดลอกเป็นศูนย์ด้วย ctypes
- 5.2 การรองรับ GQA (Grouped Query Attention) ดั้งเดิม
- 5.3 การรองรับลำดับความยาวแปรผัน
- หก การบูรณาการ vLLM: สถาปัตยกรรม Monkey-Patch ขณะรันไทม์
- 6.1 การฉีดฮุคแบบไม่รุกราน
- 6.2 เหตุใดจึงใช้ monkey-patch แทน fork
- เจ็ด การวิเคราะห์ประสิทธิภาพและการตัดสินใจทางวิศวกรรม
- 7.1 แหล่งที่มาสามชั้นของการเร่งความเร็ว
- 7.2 ค่าใช้จ่ายในการคอมไพล์ครั้งแรก
- 7.3 การรับประกันความแม่นยำแบบผสม
- 7.4 การดีบักและการสังเกตการณ์
- แปด ข้อมูลเชิงลึกทางเทคนิคและบทสรุป
unsetunsetเริ่มต้นใช้งานอย่างรวดเร็วunsetunset
โปรเจกต์นี้ออกแบบมาเป็น “ปลั๊กอิน” สามารถทดสอบประสิทธิภาพเคอร์เนลได้อย่างอิสระ หรือบูรณาการเข้ากับ vLLM-Ascend เพื่อการอนุมานแบบ end-to-end ได้อย่างราบรื่น
# โคลนที่เก็บ (รวมถึงซับโมดูล PTO-ISA)
git clone --recursive https://github.com/huawei-csl/megagdn-pto.git
cd megagdn-pto
# ติดตั้งอินเทอร์เฟซ Python (ข้อกำหนดเบื้องต้น: มีสภาพแวดล้อม CANN + torch-npu)
pip install -e '.[eval,plot]'
# การตรวจสอบความแม่นยำ
python tests/test_single_kernels.py --H-list 16,32,48,64
# การทดสอบประสิทธิภาพพื้นฐาน
python benchmarks/kernel/bench_gdn_kernels.py
--device npu:0 --n-seq 16 --l-seg 8192 --H-list 16,32,48,64
ข้อกำหนดด้านสภาพแวดล้อม: ฮาร์ดแวร์ Ascend 910B, เวอร์ชัน CANN 8.5+ และไลบรารี torch-npu แนะนำให้ใช้ vllm-ascend Docker image[1] เป็นสภาพแวดล้อมพื้นฐานที่พร้อมใช้งานทันที หากต้องการทดสอบเฉพาะเคอร์เนล PTO โดยไม่ต้องใช้ vLLM สามารถเลือกใช้ CANN Docker image[2] ที่เบากว่า หากต้องการประเมิน vLLM แบบ end-to-end ต้องรัน python vllm_patch/install_hook.py ก่อนเพื่อฉีดฮุกรันไทม์ สำหรับข้อมูลเพิ่มเติมเกี่ยวกับการเตรียมน้ำหนักโมเดลและขั้นตอนการประเมิน โปรดดู README.md[3]
unsetunsetหนึ่ง ภาพรวมสถาปัตยกรรมและปรัชญาการออกแบบunsetunset
1.1 สายการผลิตการคำนวณเจ็ดขั้นตอนของ GDN
Gated DeltaNet จะประมวลผลลำดับอินพุตเป็นชิ้นส่วน (chunk) ตามความยาวคงที่ C (ค่าเริ่มต้น 128 โทเค็น) ภายในแต่ละ chunk ขั้นตอนการคำนวณประกอบด้วยเจ็ดขั้นตอนดังต่อไปนี้:
cumsum → transpose → scaled_dot_kkt → solve_tril → wy_fast → chunk_h → chunk_o
ความหมายทางคณิตศาสตร์ของแต่ละขั้นตอนมีดังนี้:
- cumsum: คำนวณผลรวมนำหน้าภายใน chunk ของ logits ของเกตตามมิติเวลา ขั้นตอนถัดไปสามารถใช้
exp(g_sum[i] - g_sum[j])เพื่อรับค่าสัมประสิทธิ์การลดทอนเกตสะสมจากโทเค็น j ไปยังโทเค็น i ได้อย่างรวดเร็ว - transpose: จัดเรียงเลย์เอาต์ BSND ของรูปร่าง
[T, H]ใหม่เป็นเลย์เอาต์ต่อเนื่องตาม head ของรูปร่าง[H, T]การดำเนินการนี้ช่วยให้แน่ใจว่าเคอร์เนลถัดไปสามารถเข้าถึงหน่วยความจำอย่างต่อเนื่องตาม head - scaled_dot_kkt: คำนวณ
A = mask(K@K^T · gating_coeff)โดยที่สัมประสิทธิ์เกตคือexp(clamp(g[i]+log(β[i])-g[j], max=0))และซ้อนทับด้วยหน้ากากเชิงสาเหตุรูปสามเหลี่ยมล่าง ซึ่งก่อตัวเป็น “เมทริกซ์ความสนใจ” ภายใน chunk - solve_tril: ดำเนินการย้อนกลับแบบเรียกซ้ำ
A^{-1}บนเมทริกซ์สามเหลี่ยมล่าง A โดยมีวัตถุประสงค์เพื่อ “แยกส่วน” สำหรับการแยกส่วน WY ในภายหลัง - wy_fast: คำนวณสองสาขาของการแยกส่วน WY ตาม
A^{-1}:U = (A·β_2d) @ VและW = (A·(exp(g)·β)_2d) @ K - chunk_h: อัปเดตเมทริกซ์สถานะแฝง S (รูปร่าง
[H, D, D]) แบบเรียกซ้ำ และคำนวณค่าการแทรกสอดv_newพร้อมกัน - chunk_o: รวมเอาต์พุต Attention Q·K^T ภายใน chunk กับส่วนสนับสนุนสถานะข้าม chunk Q·S เพื่อให้ได้เอาต์พุตสุดท้าย O
1.2 คอขวดประสิทธิภาพของการดำเนินการแบบแยกขั้นตอน
ในโหมดการดำเนินการแบบแยกขั้นตอน แต่ละขั้นตอนต้องผ่านกระบวนการต่อไปนี้:
- เลเยอร์ Python เรียก
lib.call_kernel()ผ่าน ctypes ทริกเกอร์การจัดตารางเวลารันไทม์ของ NPU และเริ่ม AI Core ในที่สุด - ผลลัพธ์ของขั้นตอนถูกเขียนกลับไปยังหน่วยความจำส่วนกลาง (HBM)
- เมื่อขั้นตอนถัดไปเริ่มต้น จะอ่านข้อมูลเหล่านี้อีกครั้งจาก HBM
ในสถานการณ์การอนุมานทั่วไป (เช่น 128 โทเค็น × 16 head) เวลาในการคำนวณของขั้นตอนเดียวใช้เวลาเพียงไม่กี่สิบไมโครวินาที อย่างไรก็ตาม การเริ่มต้นเคอร์เนลแต่ละครั้ง (kernel launch) ต้องใช้ค่าใช้จ่ายหลายสิบถึงร้อยไมโครวินาที ความหน่วงในการจัดตารางเวลาสะสมจากการเริ่มต้นเจ็ดครั้งบวกกับความหน่วงในการอ่าน/เขียน HBM อาจคิดเป็น 30% ถึง 50% ของเวลาทั้งหมด
1.3 วิธีการแก้ปัญหาของ Megakernel
กลยุทธ์ของ Megakernel นั้นตรงไปตรงมา: คอมไพล์โค้ดฝั่งอุปกรณ์ของทั้งเจ็ดขั้นตอนลงในไฟล์ .so เดียวกัน โดยต้องดำเนินการเริ่มต้น (launch) เพียงครั้งเดียว แต่ละขั้นตอนประสานงานกันผ่านกลไกการซิงโครไนซ์ FFTS ภายในคอร์ จึงหลีกเลี่ยงการสื่อสารไปกลับระหว่าง Host และ Device หลายครั้ง ในอุดมคติ ผลลัพธ์ขั้นกลางสามารถคงอยู่ใน UB (SRAM) บนชิปได้บางส่วน ซึ่งช่วยลดจำนวนการเข้าถึง HBM ได้อย่างมาก
1.4 การจัดระเบียบโค้ดโปรเจกต์
megagdn-pto/
├── kernels/pto/ # ซอร์สโค้ดเคอร์เนล C++ PTO
│ ├── chunk_cumsum.cpp # ผลรวมนำหน้า (Vec เท่านั้น ประมาณ 430 บรรทัด)
│ ├── scaled_dot_kkt.cpp # KKT แบบมีเกต (Cube+Vec ประมาณ 700 บรรทัด)
│ ├── tri_inverse*.cpp # การย้อนกลับสามเหลี่ยม (CubeCore ประมาณ 37000 บรรทัด)
│ ├── wy_fast.cpp # การแยกส่วน WY (Vec+Cube ประมาณ 1000 บรรทัด)
│ ├── chunk_h.cpp # การเรียกซ้ำสถานะ (Cube+Vec ประมาณ 37000 บรรทัด)
│ ├── chunk_o.cpp # การคำนวณเอาต์พุต (Cube+Vec ประมาณ 60000 บรรทัด)
│ └── mega_kernel.cpp # จุดเข้าแบบรวมเจ็ดในหนึ่ง (โค้ดกาวประมาณ 500 บรรทัด)
├── megagdn_pto/ # เลเยอร์อินเทอร์เฟซ Python
│ ├── compile.py # การจัดการคอมไพล์ JIT ของ Bisheng
│ ├── kernel_libs.py # การโหลดและรันแบบแยกขั้นตอน
│ ├── mega_kernel.py # จุดเข้า Python ของเคอร์เนลแบบรวม
│ └── fast_inverse.py # การห่อหุ้ม Python ของการย้อนกลับสามเหลี่ยม
├── vllm_patch/ # monkey-patch รันไทม์ vLLM
└── third_party/pto-isa/ # ไฟล์ส่วนหัว PTO-ISA (ซับโมดูล git)
สอง โมเดลการเขียนโปรแกรมฮาร์ดแวร์ NPU: Vec, Cube และไปป์ไลน์แบบชัดแจ้ง
2.1 สถาปัตยกรรมภายใน AI Core
เพื่อทำความเข้าใจโค้ดข้างต้น ก่อนอื่นต้องเข้าใจสถาปัตยกรรม AI Core ของ Ascend 910B แต่ละ AI Core มีความคล้ายคลึงกับ SM (Stream Multiprocessor) ใน GPU โดยประมาณ แต่โครงสร้างภายในแตกต่างกันอย่างสิ้นเชิง:
| หน่วยฮาร์ดแวร์ | ฟังก์ชัน | การเปรียบเทียบกับ GPU | เลเยอร์หน่วยความจำที่ทำงาน |
|---|---|---|---|
| Vec | การดำเนินการเวกเตอร์ SIMD (add/mul/exp/log/cvt) | CUDA Core | UB (SRAM บนชิป ประมาณ 256KB) |
| Cube | การคูณเมทริกซ์ (อินพุต fp16, สะสม fp32) | Tensor Core | L0A/L0B → L0C (ไฟล์รีจิสเตอร์) |
| MTE2 | โหลด DMA: GM → UB หรือ GM → L1 | หน่วยโหลดหน่วยความจำส่วนกลาง | GM → L1/UB |
| MTE3 | จัดเก็บ DMA: UB → GM หรือ L0C → GM | หน่วยจัดเก็บหน่วยความจำส่วนกลาง | L0C/UB → GM |
| MTE1 | การลำเลียง L1 → L0A/L0B | — | L1 → L0 |
คุณสมบัติหลักคือ: ไปป์ไลน์ทั้งห้านี้ทำงานแบบขนานและอะซิงโครนัสทางกายภาพ โปรแกรมเมอร์ต้องจัดการการพึ่งพาข้อมูลระหว่างไปป์ไลน์เหล่านี้อย่างชัดเจนผ่าน set_flag/wait_flag ซึ่งเปรียบเสมือนการจัดการ memory fence ด้วยตนเองในระบบที่ไม่มีความสอดคล้องกันของแคชฮาร์ดแวร์
จุดพิเศษอีกประการคือ Cube และ Vec ทำงานบนคอร์ทางกายภาพที่แตกต่างกัน การสื่อสารระหว่างทั้งสองต้องผ่านหน่วยความจำส่วนกลาง (Global Memory) บวกกับสัญญาณข้ามคอร์ (FFTS) ซึ่งมีค่าใช้จ่ายสูงกว่าการซิงโครไนซ์ไปป์ภายในคอร์มาก
2.2 ลำดับชั้นหน่วยความจำและนามธรรม Tile
โมเดลการเขียนโปรแกรมหลักของ PTO-ISA หมุนรอบเทมเพลต
Tileโดยพื้นฐานแล้ว Tile คือตัวอธิบายพื้นที่หน่วยความจำที่ประกอบด้วยข้อมูลรูปร่าง เลย์เอาต์ และการจัดตำแหน่ง:
// ที่มา: kernels/pto/chunk_cumsum.cpp
// Tile สองมิติบน UB: ChunkSize แถว × HTC คอลัมน์, ชนิด float, เรียงตามแถว
// การจัดตำแหน่ง 512 ไบต์เป็นข้อกำหนดฮาร์ดแวร์ของ DMA
template <typename T, int R, int C, int RV = R, int CV = C,
PadValue P = PadValue::Null>
using UbND = Tile<TileType::Vec, T, R, C, BLayout::RowMajor,
RV, CV, SLayout::NoneBox, 512, P>;
เปรียบเทียบกับ PyTorch: UbND<float, 128, 16> เทียบเท่ากับ torch.empty(128, 16, dtype=float32) แต่จัดสรรใน SRAM บนชิปของแต่ละ AI Core ที่อยู่ถูกจัดการด้วยตนเองโดยโปรแกรมเมอร์ผ่าน TASSIGN(tile, byte_offset) ซึ่งเป็นพูลหน่วยความจำแบบคงที่โดยสมบูรณ์ ไม่มี malloc, ไม่มีการเก็บขยะ ทุกอย่างถูกกำหนดในช่วงคอมไพล์
GlobalTensor คือ “มุมมอง” ของเทนเซอร์ใน HBM:
// ที่มา: kernels/pto/chunk_cumsum.cpp
// เทียบเท่ากับ torch.as_strided(ptr, size=[valid, NumHeads], stride=[NumHeads, 1])
using GmShape = Shape<1, 1, 1, DYNAMIC, DYNAMIC>;
using GmStride = Stride<1, 1, 1, NumHeads, 1>;
using GmFloat = GlobalTensor<float, GmShape, GmStride>;
2.3 ภาพรวมของดั้งเดิมการซิงโครไนซ์
รูปแบบการซิงโครไนซ์ที่ปรากฏซ้ำแล้วซ้ำอีกในโค้ดนี้สามารถสรุปได้เป็นสามประเภทดังต่อไปนี้:
การซิงโครไนซ์ระหว่างไปป์ภายในคอร์: set_flag(SRC_PIPE, DST_PIPE, event) + wait_flag(...)
- ตัวอย่าง: หลังจาก MTE2 โหลด DMA เสร็จ Vec จึงจะอ่านข้อมูล UB ได้
- เปรียบเทียบ: สัญญาณระหว่างผู้ผลิต-ผู้บริโภค
สิ่งกีดขวางไปป์ทั้งหมดภายในคอร์: pipe_barrier(PIPE_ALL) หรือ pipe_barrier(PIPE_V)
- รับประกันว่าการดำเนินการที่รออยู่ทั้งหมดของไปป์ใดไปป์หนึ่งจะถูกระบายออกจนหมด
- เปรียบเทียบ: เวอร์ชันน้ำหนักเบาของ
__syncthreads()
การซิงโครไนซ์ข้ามคอร์: ffts_cross_core_sync() + wait_flag_dev(flag_id)
- ใช้สำหรับการจับมือระหว่างคอร์ Cube และคอร์ Vec
- เปรียบเทียบ: สิ่งกีดขวางระหว่างสอง rank ใน MPI
สาม ระบบคอมไพล์: การปรับแต่งเฉพาะช่วงคอมไพล์ของ Bisheng JIT
3.1 กลยุทธ์การแทรกค่าคงที่ในช่วงคอมไพล์
พารามิเตอร์ที่ไวต่อประสิทธิภาพทั้งหมด เช่น จำนวน head H, จำนวน key head Hg, มิติ head D, ขนาด chunk C จะถูกแทรกผ่านมาโคร
-Dของคอมไพเลอร์ ซึ่งหมายความว่า:
- ขอบเขตการวนซ้ำทั้งหมดทราบในช่วงคอมไพล์ คอมไพเลอร์สามารถขยายลูปได้อย่างสมบูรณ์
- ที่อยู่ UB ถูกคำนวณแบบคงที่ในช่วงคอมไพล์ ไม่จำเป็นต้องใช้เลขคณิตที่อยู่ใดๆ ในรันไทม์
- โค้ดที่สร้างจากอินสแตนซ์ของเทมเพลตคือ รูปแบบที่เหมาะสมที่สุดภายใต้ชุดพารามิเตอร์นั้น
สี่ การวิเคราะห์เคอร์เนลหลักทีละระดับ
4.1 chunk_cumsum: การดำเนินการผลรวมนำหน้าแบบเวกเตอร์ล้วน
ขั้นตอนนี้เป็นขั้นตอนที่ง่ายที่สุดในบรรดาเคอร์เนลทั้งหมด เกี่ยวข้องเฉพาะการดำเนินการ SIMD ไม่ใช้เอนจินการคำนวณเมทริกซ์เลย สาระสำคัญทางคณิตศาสตร์คือ: ดำเนินการผลรวมนำหน้าอย่างอิสระสำหรับแต่ละ head ความสนใจในแต่ละ chunk นั่นคือ g_sum[t,h] = Σ_{i=0}^{t} g[i,h]
// ที่มา: kernels/pto/chunk_cumsum.cpp (เวอร์ชันย่อของลูปหลัก)
// ขั้นตอนที่ 1: การส่ง DMA — โหลดค่าเกตของหนึ่ง chunk จากหน่วยความจำส่วนกลาง (GM) ไปยังบัฟเฟอร์แบบรวม (UB)
TLOAD(g_load, g_gm); // การดำเนินการแบบอะซิงโครนัส! ดำเนินการบนไปป์ไลน์ MTE2
// ขั้นตอนที่ 2: รอการซิงโครไนซ์ — ตรวจสอบให้แน่ใจว่าการส่ง DMA เสร็จสมบูรณ์
set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
// ขั้นตอนที่ 3: การคำนวณผลรวมนำหน้า — head ทั้งหมด H ดำเนินการแบบขนานภายในความกว้าง SIMD
UbND<float, 1, HTC> acc_ub; // ตัวสะสม ความกว้างครอบคลุมทุก head
TASSIGN(acc_ub, AccUbAddr);
TMOV(acc_ub, g_row_0); // เริ่มต้น: acc = g[0, :]
pipe_barrier(PIPE_V);
for (int32_t i = 1; i < valid; ++i) {
UbND<float, 1, HTC> g_row_i;
TASSIGN(g_row_i, GUbAddr + i * RowBytes);
TADD(acc_ub, acc_ub, g_row_i); // สะสมแบบขนานทุก head: acc += g[i, :]
pipe_barrier(PIPE_V);
UbND<float, 1, HTC> s_row_i;
TASSIGN(s_row_i, SUbAddr + i * RowBytes);
TMOV(s_row_i, acc_ub); // เขียนผลลัพธ์: g_sum[i, :] = acc
pipe_barrier(PIPE_V);
}
// ขั้นตอนที่ 4: รอการซิงโครไนซ์ + เขียนกลับ DMA
set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
TSTORE(gs_gm, s_store);
กลไกการกระจายงาน ใช้รูปแบบที่คล้ายกับ grid-stride loop ของ CUDA: แต่ละ AI Core ประมวลผล chunk ที่แตกต่างกันในลักษณะแบบ round-robin HTC = ((NumHeads + 7) / 8) * 8 ในโค้ดจะปรับจำนวน head ให้เป็นพหุคูณของ 8 (นั่นคือการจัดตำแหน่ง 32 ไบต์) เพื่อให้แน่ใจว่า “คำ” SIMD แต่ละคำถูกเติมเต็มอย่างสมบูรณ์ หลีกเลี่ยงการสิ้นเปลืองทรัพยากร
สิ่งที่ควรทราบเป็นพิเศษคือ pipe_barrier(PIPE_V) ที่ดูเหมือนซ้ำซ้อนในโค้ดนั้นจำเป็นจริงๆ — เอนจิน Vec ใช้สถาปัตยกรรมไปป์ไลน์ ผลลัพธ์การคำนวณของคำสั่งก่อนหน้าอาจยังไม่ได้เขียนไปยัง UB ในขณะที่คำสั่งถัดไปเริ่มอ่านข้อมูลแล้ว เฉพาะการดำเนินการ barrier อย่างชัดแจ้งเท่านั้นที่รับประกันลำดับการดำเนินการที่ถูกต้อง
4.2 scaled_dot_kkt: ความร่วมมือแบบดูอัลคอร์ของ Cube และ Vec
นี่เป็นเคอร์เนลที่มีความซับซ้อนมากที่สุดในบรรดาเคอร์เนลแบบขั้นตอนเดียวทั้งหมด มันระดมเอนจิน Cube เพื่อดำเนินการคูณเมทริกซ์ และเอนจิน Vec เพื่อคำนวณสัมประสิทธิ์เกตพร้อมกัน เอนจินทั้งสองทำงานบนคอร์ทางกายภาพที่แตกต่างกัน โดยประสานงานกันผ่านพื้นที่ทำงาน (GM) และกลไกสัญญาณ FFTS
ฝั่ง Cube (การคูณเมทริกซ์ K@K^T)
// ที่มา: kernels/pto/scaled_dot_kkt.cpp (ส่วนที่เกี่ยวข้องกับ Cube, ย่อ)
#if defined(__DAV_C220_CUBE__)
// โหลด K [C×D] จากหน่วยความจำส่วนกลาง (GM) ไปยังแคชระดับหนึ่ง (L1) ในรูปแบบ NZ fractal
TLOAD(_l1, _gm);
// หากข้อมูลท้ายไม่ถึง ChunkSize ให้เติมศูนย์
if (valid_rows != ChunkSize) TFILLPAD(_l1, _l1);
// เทคนิคทรานสโพส: TRESHAPE คือการดำเนินการตีความเลย์เอาต์ใหม่โดยไม่มีค่าใช้จ่าย
// โดยการแปลง NZ เป็น ZN จะได้ K^T ในเชิงตรรกะ โดยไม่ต้องย้ายข้อมูลจริง!
L1MatZN<half, HiddenSize, ChunkSize> _bzn;
TRESHAPE(_bzn, k_l1); // ได้ K^T ฟรี
// ใช้ไปป์ไลน์ MTE1 ลำเลียงข้อมูลจาก L1 ไปยัง L0A/L0B
TEXTRACT(_l0a, k_l1, 0, 0); // ตัวถูกดำเนินการซ้าย: K
TEXTRACT(_l0b, _bzn, 0, 0); // ตัวถูกดำเนินการขวา: K^T
// ดำเนินการคูณเมทริกซ์: K @ K^T สะสมผลลัพธ์ไปยัง L0C (fp16×fp16→fp32)
TMATMUL(a_l0, _l0a, _l0b);
// เขียนผลลัพธ์ L0C กลับไปยัง workspace หน่วยความจำส่วนกลาง และตัดทอน fp32→fp16 โดยอัตโนมัติ
TSTORE(_gm, _l0);
// สัญญาณซิงโครไนซ์ข้ามคอร์: แจ้งเอนจิน Vec
⚠️ หมายเหตุ: เนื้อหาได้รับการแปลโดย AI และตรวจสอบโดยมนุษย์ หากมีข้อผิดพลาดโปรดแจ้ง
☕ สนับสนุนค่ากาแฟทีมงาน
หากคุณชอบบทความนี้ สามารถสนับสนุนเราได้ผ่าน PromptPay
SCAN TO PAY WITH ANY BANK
本文来自网络搜集,不代表คลื่นสร้างอนาคต立场,如有侵权,联系删除。转载请注明出处:https://www.itsolotime.com/th/archives/35430
