Gated DeltaNet เพิ่มความเร็วในการอนุมาน 3 เท่า! Huawei ใช้ PTO-ISA เขียน Megakernel ด้วยมือ ไล่ตีเส้นฐาน Triton จนยับ

ในสนามประลองการอนุมานของโมเดลภาษาขนาดใหญ่ ประสิทธิภาพการคำนวณของกลไก 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 คอขวดประสิทธิภาพของการดำเนินการแบบแยกขั้นตอน

ในโหมดการดำเนินการแบบแยกขั้นตอน แต่ละขั้นตอนต้องผ่านกระบวนการต่อไปนี้:

  1. เลเยอร์ Python เรียก lib.call_kernel() ผ่าน ctypes ทริกเกอร์การจัดตารางเวลารันไทม์ของ NPU และเริ่ม AI Core ในที่สุด
  2. ผลลัพธ์ของขั้นตอนถูกเขียนกลับไปยังหน่วยความจำส่วนกลาง (HBM)
  3. เมื่อขั้นตอนถัดไปเริ่มต้น จะอ่านข้อมูลเหล่านี้อีกครั้งจาก 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

PromptPay QR
SCAN TO PAY WITH ANY BANK

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

Like (0)
Previous 3 hours ago
Next 3 hours ago

相关推荐