1. คลื่นสร้างอนาคตHome
  2. โครงการโอเพนซอร์ส

บอกลาปัญหาโอเปอเรเตอร์ระดับล่าง: Hugging Face Kernel Hub ทำให้ CUDA และ Metal Kernel ใช้งานได้แบบปลั๊กแอนด์เพลย์เหมือนโมเดล เพิ่มความเร็วในการอนุมานได้ทันที

• โครงการโอเพนซอร์ส • 14 views

บอกลาปัญหาโอเปอเรเตอร์ระดับล่าง: Kernel Hub จาก Hugging Face ทำให้ CUDA และ Metal Kernels เสียบปลั๊กและใช้งานได้เหมือนโมเดล การเร่งความเร็วอนุมานอยู่แค่เอื้อม

ในระบบโมเดลขนาดใหญ่ คอขวดด้านประสิทธิภาพมักไม่ได้มาจากสถาปัตยกรรมของโมเดลเอง แต่ซ่อนอยู่ในโอเปอเรเตอร์ระดับล่างที่ดูเหมือนไม่มีนัยสำคัญ การดำเนินการแบบฟิวชันของฟังก์ชันกระตุ้น การเข้าถึง KV Cache การควอนไทซ์และดีควอนไทซ์น้ำหนัก 4-bit รายละเอียดเหล่านี้เป็นตัวกำหนดว่า GPU หรือ Apple Silicon จะสามารถดึงศักยภาพทั้งหมดออกมาได้จริงหรือไม่

ในอดีต เคอร์เนลการคำนวณประสิทธิภาพสูงเหล่านี้กระจัดกระจายอยู่ในโปรเจกต์ต่างๆ เช่น vLLM, FlashAttention, bitsandbytes, MLX, Triton ทำให้ขั้นตอนการติดตั้งซับซ้อน สภาพแวดล้อมแบ็กเอนด์แยกส่วน และการจัดการเวอร์ชันทำได้ยาก

โปรเจกต์ huggingface/kernels-community พยายามตอบคำถามพื้นฐานกว่านั้น: เราสามารถดึงเคอร์เนลการคำนวณระดับล่างมาใช้ตามความต้องการได้เหมือนกับการโหลดโมเดลจาก Hub ได้หรือไม่?

นี่ไม่ใช่ไลบรารีอัลกอริทึมเดี่ยวๆ แต่เป็นคลังเก็บซอร์สโค้ดเคอร์เนลที่ผสานรวม CUDA, Metal, Triton, การผูก C++/Python เข้ากับระบบการแจกจ่ายของ Hugging Face ได้อย่างราบรื่น

  • Kernel Hub อนุญาตให้ไลบรารี Python และแอปพลิเคชันต่างๆ โหลด Kernel การคำนวณที่ปรับแต่งแล้วจาก Hugging Face Hub ได้โดยตรง คุณสามารถมองว่ามันเป็นแพลตฟอร์มที่คล้ายกับ Model Hub แต่เชี่ยวชาญในการจัดเก็บโค้ดประสิทธิภาพสูงระดับล่าง (Kernel) ซึ่งโดยทั่วไปแล้วสามารถเร่งการดำเนินการเฉพาะต่างๆ บนอุปกรณ์กราฟิกได้
  • ไม่จำเป็นต้องจัดการกับสภาพแวดล้อม dependencies ที่ซับซ้อน พารามิเตอร์การคอมไพล์ที่ต้องดีบัก หรือคอมไพล์ไลบรารีโปรแกรม เช่น Triton, CUTLASS จากซอร์สโค้ดอีกต่อไป ด้วยความช่วยเหลือของ Kernel Library คุณสามารถรับและรันเคอร์เนลการคำนวณที่คอมไพล์ไว้ล่วงหน้าและปรับแต่งประสิทธิภาพแล้วได้ตามต้องการ

คลังโค้ด: https://github.com/huggingface/kernels-community
เอกสารทางการ: https://huggingface.co/docs/kernels/

หนึ่ง เริ่มต้นใช้งานอย่างรวดเร็ว: โหลดเคอร์เนลเหมือนโหลดโมเดล

README ในไดเรกทอรีรากของ kernels-community อธิบายตำแหน่งของโปรเจกต์อย่างระมัดระวัง: Kernel Hub อนุญาตให้ไลบรารี Python และแอปพลิเคชันโหลดเคอร์เนลการคำนวณจาก Hugging Face Hub ได้โดยตรง คลังเก็บนี้จะเก็บซอร์สโค้ดที่เผยแพร่ไปยัง hf.co/kernels-community[1] เอกสาร Kernel Hub ที่สมบูรณ์ยิ่งขึ้นสามารถอ้างอิงได้ที่ huggingface.co/docs/kernels/index[2]

สารบัญบทความนี้

  • หนึ่ง เริ่มต้นใช้งานอย่างรวดเร็ว: โหลดเคอร์เนลเหมือนโหลดโมเดล
  • สอง ปัญหาที่โปรเจกต์แก้ไขได้จริง: การทำให้เคอร์เนลประสิทธิภาพสูงเป็น “Hub”
    • 2.1 จาก “ซอร์สโค้ดที่กระจัดกระจาย” สู่ “โหลดได้ตามต้องการ”
    • 2.2 ไดเรกทอรีคลังเก็บคือแผนที่ประสิทธิภาพของระบบ AI
  • สาม รูปแบบการห่อหุ้มแบบรวม: README, build.toml, flake.nix และ torch-ext
    • 3.1 โครงสร้างมาตรฐานของแพ็คเกจเคอร์เนล
    • 3.2 flake.nix: มอบหมายข้อตกลงการสร้างให้กับ kernel-builder
  • สี่ จาก Python สู่ฮาร์ดแวร์: สายการเรียกใช้ที่สมบูรณ์ของ activation kernel
    • 4.1 ชั้น Python: รักษาประสบการณ์การใช้งานที่คล้ายกับโมดูลทั่วไป
    • 4.2 ชั้นการลงทะเบียน C++: เนมสเปซของ Torch op และการจัดส่งแบ็กเอนด์
    • 4.3 ชั้น CUDA: ฟิวชันการกระตุ้นและการคูณในการเปิดตัวเคอร์เนลครั้งเดียว
    • 4.4 ชั้น Metal: MPS ไม่ใช่ “เวอร์ชันแปลของ CUDA”
  • ห้า Paged Attention: ทำให้ KV Cache กลายเป็นหน่วยความจำแบบแบ่งหน้า
    • 5.1 ทำไม Paged Attention จึงสำคัญต่อการอนุมาน
    • 5.2 เวลาคอมไพล์ การสร้างอินสแตนซ์เทมเพลต และรูปทรงที่จำกัด
  • หก bitsandbytes-mps: ทำให้การควอนไทซ์ 4-bit ทำงานบน Apple Silicon
    • 6.1 แกนหลักของ NF4/FP4 ไม่ใช่ “การบีบอัด” แต่เป็นการถอดรหัสไปพร้อมกับการคำนวณ
    • 6.2 Codebook และรูปแบบการบรรจุ: “พจนานุกรม” ของค่าบิตต่ำ
    • 6.3 BnBQuantizedBlockLoader: ซ่อนการดีควอนไทซ์ไว้ในการคูณเมทริกซ์
  • เจ็ด Triton kernel: เขียนโปรแกรมเฉพาะทางที่ใกล้เคียงฮาร์ดแวร์ด้วย Python
    • 7.1 ลักษณะทางวิศวกรรมของ gpt-oss-triton-kernels
    • 7.2 specialize.py: สร้างฟังก์ชัน Triton เฉพาะทางแบบไดนามิก
    • 7.3 MoE routing: การเรียงลำดับที่เสถียรและ scatter/gather ภายในเคอร์เนล
  • แปด RMSNorm และ CPU/XPU: ประสิทธิภาพสูงไม่ใช่เพียงแค่ของ CUDA
    • 8.1 การจัดส่ง CPU SIMD: เลือก AVX512/AVX2/Fallback ในขณะรันไทม์
    • 8.2 การห่อหุ้ม Python autograd: ทั้ง forward/backward เป็นทรัพย์สินของเคอร์เนล
  • เก้า การแบ่งงานภาษา: ทำไมคลังเก็บนี้ถึงต้องการ C++, Python, CUDA, Metal พร้อมกัน
    • 9.1 Python: API, Triton JIT และจุดเริ่มต้นการทดสอบ
    • 9.2 C++: ส่วนขยาย PyTorch และกาวสำหรับรันไทม์
    • 9.3 CUDA และ Metal: สถานที่ที่ใกล้ชิดกับฮาร์ดแวร์อย่างแท้จริง
  • ความหมายพื้นฐานของคลังเก็บนี้: ทำให้เคอร์เนลเป็นพลเมืองชั้นหนึ่งของระบบนิเวศ AI

2.1 สามขั้นตอน: การใช้ Kernel Hub ตั้งแต่เริ่มต้น

สำหรับผู้ใช้ส่วนใหญ่ เส้นทางการเริ่มต้นใช้งานนั้นง่ายมาก เพียงสามขั้นตอน: ติดตั้งไลบรารี kernels ก่อน จากนั้นใช้ฟังก์ชัน get_kernel() ดึงแพ็คเกจเคอร์เนลที่ระบุจากระยะไกล สุดท้ายก็สามารถเรียกใช้ฟังก์ชันภายในได้โดยตรงเหมือนกับโมดูล Python ทั่วไป

โค้ดด้านล่างนี้แสดงให้เห็นกระบวนการนี้อย่างชัดเจน โดยดึงเคอร์เนลฟังก์ชันกระตุ้นจากคลังเก็บ kernels-community/activation และดำเนินการ silu_and_mul บนอุปกรณ์ CUDA:

# ที่มา: activation/scripts/readme_example.py  
import torch
from kernels import get_kernel

torch.manual_seed(42)
activation = get_kernel("kernels-community/activation")
device = torch.device("cuda")

num_tokens, hidden_dim = 128, 512
input_tensor = torch.randn(
num_tokens, 2 * hidden_dim, device=device, dtype=torch.float16
)

out_shape = input_tensor.shape[:-1] + (hidden_dim,)
out_kernel = torch.empty(out_shape, dtype=input_tensor.dtype, device=device)
out_kernel = activation.silu_and_mul(out_kernel, input_tensor)

print(out_kernel.shape)  # torch.Size([128, 512])

สำหรับนักพัฒนาที่ต้องการพัฒนา或สร้างเคอร์เนลในเครื่อง การดำเนินการทั่วไปจะแสดงอยู่ใน README ของแต่ละโปรเจกต์ย่อย ตัวอย่างเช่น bitsandbytes-mps คำสั่งสร้างมีดังนี้:

# ที่มา: bitsandbytes-mps/README.md  
pip install kernel-builder
kernel-builder build .

สำหรับรายละเอียดเพิ่มเติมเกี่ยวกับการเขียนและการสร้างเคอร์เนล เอกสารการมีส่วนร่วมของคลังเก็บแนะนำให้นักพัฒนาอ้างอิงคำแนะนำเกี่ยวกับ “writing kernels” ในโปรเจกต์ kernel-builder ของ Hugging Face และเอกสารการสร้าง Nix

  • สำหรับผู้ใช้ทั่วไป: จุดเริ่มต้นที่ดีที่สุดคือการอ่านส่วน Usage ใน README ของเคอร์เนลแต่ละตัวโดยตรง
  • สำหรับผู้ดูแลโปรเจกต์: จำเป็นต้องเข้าใจข้อตกลงระหว่าง build.toml, flake.nix, torch-ext และซอร์สโค้ดแบ็กเอนด์อย่างลึกซึ้ง

สอง ปัญหาที่โปรเจกต์แก้ไขได้จริง: การทำให้เคอร์เนลประสิทธิภาพสูงเป็น “Hub”

2.1 จาก “ซอร์สโค้ดที่กระจัดกระจาย” สู่ “โหลดได้ตามต้องการ”

เคอร์เนลประสิทธิภาพสูงในระบบ AI มักเผชิญกับปัญหาหลักสามประการ

  • ประการแรก แหล่งที่มาของซอร์สโค้ดกระจัดกระจายสูง ตัวอย่างเช่น activation kernel มีต้นกำเนิดจาก vLLM, ซีรีส์ FlashAttention มาจากระบบนิเวศ Dao-AILab ในขณะที่ bitsandbytes MPS quantization kernel ยืมการออกแบบจาก MLX และ bitsandbytes แต่ละตัวมีความยอดเยี่ยมในตัวเอง แต่ถูกผูกติดกับโปรเจกต์ ระบบสร้าง และจังหวะการเผยแพร่ที่แตกต่างกัน ทำให้จัดการแบบรวมศูนย์ได้ยาก
  • ประการที่สอง แบ็กเอนด์ฮาร์ดแวร์แยกส่วนจากกัน รูปแบบการพัฒนาของ CUDA, ROCm, XPU, MPS/Metal แตกต่างกันอย่างมาก: เคอร์เนล CUDA มักเน้น C++ และการสร้างอินสแตนซ์เทมเพลต; Metal ต้องใช้การเชื่อมต่อ Objective-C++ กับ MPS command encoder; Triton มีอยู่ในรูปแบบ Python JIT การทำให้พวกมันเปิดเผยต่อแอปพลิเคชัน Python ชั้นบนอย่างเป็นหนึ่งเดียวนั้นเป็นปัญหาทางวิศวกรรมที่ใหญ่หลวง
  • ประการที่สาม ห่วงโซ่การปรับใช้ยาวเกินไป ส่วนขยาย PyTorch C++/CUDA แบบดั้งเดิมมักต้องการให้ผู้ใช้คอมไพล์ในเครื่อง ซึ่งในกระบวนการนั้นมักจะล้มเหลวได้ง่ายเนื่องจากเวอร์ชัน CUDA, คอมไพเลอร์, ABI หรือสถาปัตยกรรม GPU ที่ไม่ตรงกัน คุณค่าหลักของ Kernel Hub คือการทำให้แอปพลิเคชันสามารถแยกวิเคราะห์ ดาวน์โหลด และโหลดทรัพย์สินเคอร์เนลที่เหมาะสมจาก Hub ได้โดยตรง ซึ่งจะช่วยลดความซับซ้อนส่วนใหญ่ไว้ที่ฝั่งการสร้างและการเผยแพร่

ดังนั้น เป้าหมายหลักของ kernels-community ไม่ใช่ “การสร้างเคอร์เนล” แต่เป็นการสร้างวิธีการจัดระเบียบแบบใหม่: ทำให้เคอร์เนลแต่ละตัวกลายเป็นหน่วยที่สามารถสร้าง ทดสอบ เผยแพร่ และโหลดจาก Hub ได้ตามต้องการ

2.2 ไดเรกทอรีคลังเก็บ: แผนที่ประสิทธิภาพของระบบ AI

ในไดเรกทอรีราก คุณจะเห็นไดเรกทอรีย่อยจำนวนมากที่ตั้งชื่อตามโอเปอเรเตอร์หรืออัลกอริทึม เช่น: activation, rmsnorm, paged-attention, flash-attn2/3/4, flash-mla, mamba-ssm, quantization-bitsandbytes, bitsandbytes-mps, gpt-oss-triton-kernels, gpt-oss-metal-kernels, liger-kernels, triton-kernels เป็นต้น

นี่ไม่ใช่การกองรวมกันแบบสุ่ม แต่ครอบคลุมเส้นทางสำคัญเกือบทั้งหมดในกระบวนการอนุมานและฝึกอบรมโมเดลขนาดใหญ่สมัยใหม่:

  • การเร่งความเร็ว Attention: FlashAttention, PagedAttention, MLA;
  • การฟิวชัน Normalization และ Activation: RMSNorm, LayerNorm, SwiGLU, GeGLU;
  • การควอนไทซ์และการคำนวณบิตต่ำ: GPTQ, EETQ, bitsandbytes, FP8, NF4/FP4;
  • MoE และ Routing: scattermoe, sonic-moe, megablocks, gpt-oss Triton routing;
  • การรองรับฮาร์ดแวร์หลายตัว: CUDA, Metal/MPS, XPU, CPU SIMD, Triton

กล่าวอีกนัยหนึ่ง คลังเก็บนี้เปรียบเสมือน “ชั้นวางโอเปอเรเตอร์ระดับล่าง” Transformers, Diffusers, บริการอนุมาน หรือโค้ดวิจัยระดับบน ไม่จำเป็นต้องสร้างวงล้อขึ้นมาใหม่ แต่สามารถเลือกเคอร์เนลที่เหมาะสมที่สุดจากชั้นวางได้ตามต้องการ

สาม รูปแบบการห่อหุ้มแบบรวม: README, build.toml, flake.nix และ torch-ext

3.1 โครงสร้างมาตรฐานของแพ็คเกจเคอร์เนล

เอกสารการมีส่วนร่วมอธิบายขั้นตอนมาตรฐานสำหรับการเพิ่มเคอร์เนลใหม่: ขั้นแรกสร้างไดเรกทอรี จากนั้นเพิ่ม README.md, flake.nix, build.toml ตามลำดับ และวางซอร์สโค้ดในไดเรกทอรีที่กำหนด หากเคอร์เนลนั้นไม่ใช่การใช้งาน Triton ล้วนๆ จำเป็นต้องเพิ่ม torch-ext และ torch_binding.cpp เพิ่มเติม เพื่อลงทะเบียนโอเปอเรเตอร์ Torch ผ่านกลไกส่วนขยาย PyTorch

ในที่สุดสิ่งนี้ก็ก่อให้เกิดโครงสร้างแบบชั้นที่ชัดเจน:

  • README.md: สำหรับผู้ใช้ อธิบายที่มา ความสามารถ การใช้งาน และผลการทดสอบประสิทธิภาพของเคอร์เนล
  • build.toml: สำหรับตัวสร้าง ประกาศชื่อ เวอร์ชัน ใบอนุญาต แบ็กเอนด์ที่รองรับ และตำแหน่งซอร์สโค้ดของเคอร์เนล
  • flake.nix: สำหรับ Nix และ kernel-builder ใช้สร้างเอาต์พุตการสร้างที่ทำซ้ำได้
  • torch-ext: สำหรับ Python/PyTorch ให้จุดเข้าโมดูลและการผูกโอเปอเรเตอร์ที่กำหนดเอง
  • ไดเรกทอรีซอร์สโค้ดแบ็กเอนด์: สำหรับฮาร์ดแวร์ เก็บการใช้งานสำหรับแบ็กเอนด์ต่างๆ เช่น CUDA, Metal, C++, SYCL, Triton

ตัวอย่างเช่น activation/build.toml ไฟล์การกำหนดค่านี้ระบุอย่างชัดเจนว่า activation kernel เดียวกันรองรับทั้งแบ็กเอนด์ CUDA และ Metal พร้อมกัน:

# ที่มา: activation/build.toml  
[general]  
name = "activation"  
version = 1  
license = "Apache-2.0"  
backends = ["cuda", "metal"]  

[general.hub]  
repo-id = "kernels-community/activation"  

[torch]  
src = [  
"torch-ext/torch_binding.cpp",  
"torch-ext/torch_binding.h",  
]  

[kernel.activation_metal]  
backend = "metal"  
depends = ["torch"]  
src = [  
"activation_metal/activation.mm",  
"activation_metal/activation.metal",  
]  

[kernel.activation]  
backend = "cuda"  
depends = ["torch"]  
src = [  
"activation/activation_kernels.cu",  
"activation/cuda_compat.h",  
"activation/dispatch_utils.h",  
]

บทบาทของการกำหนดค่านี้มีความสำคัญมาก: มันแยก “แพ็คเกจเคอร์เนลแบบรวมในสายตาผู้ใช้” ออกเป็นการใช้งานอิสระสำหรับหลายแบ็กเอนด์ จากมุมมองการเรียกใช้ระดับบน ยังคงเป็น activation.silu_and_mul() แต่ในขั้นตอนการสร้างและการโหลด ระบบจะเลือกผลิตภัณฑ์การสร้างของ CUDA หรือ Metal โดยอัตโนมัติตามแพลตฟอร์มที่รันอยู่

3.2 flake.nix: มอบหมายข้อตกลงการสร้างให้กับ kernel-builder

ไฟล์ flake.nix ของโปรเจกต์ย่อยส่วนใหญ่นั้นเรียบง่ายมาก ตรรกะหลักคือการนำเข้า kernel-builder และเรียกใช้ genKernelFlakeOutputs:

# ที่มา: activation/flake.nix  
{  
description = "Flake for activation kernels";  

inputs = {  
kernel-builder.url = "github:huggingface/kernels/torch-2.12";  
};  

outputs = { self, kernel-builder }:  
kernel-builder.lib.genKernelFlakeOutputs {  
inherit self;  
path = ./.;  
};  
}

การออกแบบนี้บ่งชี้ว่า kernels-community ไม่ได้เขียนตรรกะการสร้างซ้ำในแต่ละไดเรกทอรีย่อย แต่รวมกฎการสร้างไว้ใน kernel-builder โปรเจกต์ย่อยเพียงแค่ต้องบอกว่า “ฉันคือใคร มีซอร์สโค้ดอะไรบ้าง รองรับแบ็กเอนด์ใดบ้าง” ส่วนกระบวนการสร้าง บรรจุ อัปโหลดไปยัง Hub ที่เฉพาะเจาะจงนั้นทั้งหมดจะถูกจัดการโดยชุดเครื่องมือแบบรวม

unsetunsetสี่ จาก Python สู่ฮาร์ดแวร์: สายการเรียกใช้ที่สมบูรณ์ของ activation kernelunsetunset

4.1 ชั้น Python: รักษาประสบการณ์การใช้งานที่คล้ายกับโมดูลทั่วไป

จุดเข้าแพ็คเกจ Python ของ activation นั้นเบามาก โดยเพียงแค่ส่งต่อฟังก์ชันไปยัง _ops:

# ที่มา: activation/torch-ext/activation/__init__.py  
import torch  
from ._ops import ops  

def silu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:  
ops.silu_and_mul(out, x)  
return out  

def gelu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:  
ops.gelu_and_mul(out, x)  
return out

แม้การออกแบบนี้จะดูเรียบง่าย แต่ก็มีความสำคัญอย่างยิ่ง สำหรับนักพัฒนาแอปพลิเคชัน มันซ่อนรายละเอียดระดับล่างของ C++ extension สำหรับผู้ดูแลเคอร์เนล มันเป็นชั้น API ที่เสถียร ซึ่งอนุญาตให้เปลี่ยน CUDA, Metal หรือแบ็กเอนด์อื่นๆ ภายในได้อย่างยืดหยุ่น

layers.py ห่อหุ้มฟังก์ชันเหล่านี้เป็น nn.Module ต่อไป เช่น SiluAndMul จะตรวจสอบความต่อเนื่อง จัดสรรเทนเซอร์เอาต์พุต จากนั้นเรียกใช้ op ที่กำหนดเอง ทำให้เคอร์เนลสามารถฝังตัวในโครงสร้างโมเดลได้อย่างเป็นธรรมชาติมากขึ้น

4.2 ชั้นการลงทะเบียน C++: เนมสเปซของ Torch op และการจัดส่งแบ็กเอนด์

กุญแจสำคัญที่เชื่อมต่อการเรียกใช้ Python เข้ากับการใช้งานฮาร์ดแวร์จริงๆ อยู่ที่ torch_binding.cpp มันกำหนด schema ของ op ผ่าน TORCH_LIBRARY_EXPAND และเลือกที่จะลงทะเบียนการใช้งาน CUDA หรือ MPS ตามมาโครการคอมไพล์:

// ที่มา: activation/torch-ext/torch_binding.cpp  
#include <torch/library.h>  
#include "registration.h"  
#include "torch_binding.h"  

TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {  
ops.def("silu_and_mul(Tensor! out, Tensor input) -> ()");  
#if defined(CUDA_KERNEL)  
ops.impl("silu_and_mul", torch::kCUDA, &silu_and_mul);  
#elif defined(METAL_KERNEL)  
ops.impl("silu_and_mul", torch::kMPS, &silu_and_mul);  
#endif  

ops.def("gelu_and_mul(Tensor! out, Tensor input) -> ()");  
#if defined(CUDA_KERNEL)  
ops.impl("gelu_and_mul", torch::kCUDA, &gelu_and_mul);  
#elif defined(METAL_KERNEL)  
ops.impl("gelu_and_mul", torch::kMPS, &gelu_and_mul);  
#endif  
}

คุณสามารถมองชั้นนี้เป็น “ด่านศุลกากรโอเปอเรเตอร์”: ฝั่ง Python รู้เพียง ops.silu_and_mul แต่ชั้นการลงทะเบียน C++ จะกำหนดเส้นทางการเรียกไปยัง CUDA kernel หรือ Metal kernel ตามอุปกรณ์ที่เทนเซอร์อยู่และแบ็กเอนด์ที่สร้างขึ้น

4.3 ชั้น CUDA: ฟิวชันการกระตุ้นและการคูณในการเปิดตัวเคอร์เนลครั้งเดียว

นิพจน์ทางคณิตศาสตร์ของ SwiGLU มีดังนี้:

input = [x1, x2]  
output = silu(x1) * x2

หากใช้การดำเนินการแบบผสมผสานของ PyTorch ทั่วไป มักจะต้องดำเนินการหลายขั้นตอนอิสระ เช่น การแบ่งส่วน การกระตุ้น และการคูณตามลำดับ ในขณะที่ fused kernel สามารถอ่านข้อมูลครั้งเดียวและคำนวณทั้งหมดเสร็จสิ้น ซึ่งช่วยหลีกเลี่ยงค่าใช้จ่ายเพิ่มเติมจากการสร้างเทนเซอร์กลางและการเปิดตัวเคอร์เนลได้อย่างมีประสิทธิภาพ

// ที่มา: activation/activation/activation_kernels.cu  
#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, ACT_FIRST)                      
int d = input.size(-1) / 2;                                                
int64_t num_tokens = input.numel() / input.size(-1);                       
dim3 grid(num_tokens);                                                     
dim3 block(std::min(d, 1024));                                             
if (num_tokens == 0) { return; }                                           
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();              
VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "act_and_mul_kernel", [&] {   
vllm::act_and_mul_kernel<scalar_t, KERNEL<scalar_t>, ACT_FIRST>          
<<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(),                 
input.data_ptr<scalar_t>(), d);           
});  

void silu_and_mul(torch::Tensor& out, torch::Tensor& input) {  
TORCH_CHECK(input.is_contiguous());  
TORCH_CHECK(out.is_contiguous());  
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, true);  
}

โค้ดนี้สะท้อนให้เห็นถึงการแลกเปลี่ยนและการออกแบบทางวิศวกรรมหลายประการ:

  • d = input.size(-1) / 2: กำหนดให้ครึ่งแรกของมิติสุดท้ายใช้สำหรับการคำนวณกระตุ้น และครึ่งหลังใช้สำหรับการคูณแบบเกต
  • grid(num_tokens): แต่ละ token หรือหนึ่งแถวหลังจากทำให้แบนราบ สอดคล้องกับหนึ่ง block
  • block(std::min(d, 1024)): แต่ละแถวสามารถเปิดใช้งานได้สูงสุด 1024 เธรด
  • VLLM_DISPATCH_FLOATING_TYPES: สร้างการใช้งานที่สอดคล้องกันสำหรับประเภทข้อมูลทศนิยมที่แตกต่างกัน
  • getCurrentCUDAStream(): ปฏิบัติตามการจัดการสตรีมปัจจุบันของ PyTorch อย่างเคร่งครัด เพื่อหลีกเลี่ยงการทำลายความหมายของการดำเนินการแบบอะซิงโครนัส

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

4.4 ชั้น Metal: MPS ไม่ใช่ “การแปล CUDA อย่างง่าย”

activation kernel เดียวกันยังรองรับแพลตฟอร์ม Apple MPS แบ็กเอนด์ Metal ไม่ใช่การแปลโค้ด CUDA อย่างกลไก แต่จำเป็นต้องใช้ Objective-C++ เพื่อรับ MTL buffer, กำหนดค่า pipeline, จัดการเธรด และส่งงานไปยัง MPS stream

// ที่มา: activation/activation_metal/activation.mm  
static void checkInputs(torch::Tensor &out, torch::Tensor const &input) {  
TORCH_CHECK(input.device().is_mps(), "input must be a MPS tensor");  
TORCH_CHECK(input.is_contiguous(), "input must be contiguous");  
TORCH_CHECK(out.device().is_mps(), "output must be a MPS tensor");  
TORCH_CHECK(out.is_contiguous(), "output must be contiguous");  
TORCH_CHECK(input.scalar_type() == torch::kFloat ||  
input.scalar_type() == torch::kHalf,  
"Unsupported data type: ", input.scalar_type());  
}  

void silu_and_mul(torch::Tensor &out, torch::Tensor &input) {  
checkInputs(out, input);  
dispatchGatedKernel("silu_and_mul", out, input);  
}

โค้ดนี้เผยให้เห็นหน้าที่สำคัญที่สุดของแบ็กเอนด์ Metal: จุดเน้นไม่ได้อยู่ที่การเขียนสูตรคณิตศาสตร์ใหม่ แต่อยู่ที่การตรวจสอบให้แน่ใจว่าประเภทอุปกรณ์ของเทนเซอร์, dtype, เงื่อนไข contiguous และชื่อ pipeline สอดคล้องกับ Metal shader จริงอย่างเคร่งครัด สำหรับผู้ใช้ มันยังคงเป็นฟังก์ชัน Python เท่านั้น แต่สำหรับผู้ดูแล เบื้องหลังนี้คือสแต็กฮาร์ดแวร์ที่แตกต่างอย่างสิ้นเชิง

ห้า Paged Attention: ทำให้ KV Cache กลายเป็นหน่วยความจำแบบแบ่งหน้า

5.1 ความสำคัญของ Paged Attention ต่อการอนุมาน

ในกระบวนการอนุมานแบบ autoregressive ของโมเดลขนาดใหญ่ KV Cache จะเพิ่มขึ้นอย่างต่อเนื่องตามความยาวของการสร้าง

วิธีการจัดสรรหน่วยความจำแบบต่อเนื่องแบบดั้งเดิมมีแนวโน้มที่จะทำให้เกิดการแตกกระจายและสิ้นเปลืองทรัพยากร โดยเฉพาะอย่างยิ่งในสถานการณ์ที่มีหลายคำขอ หลาย batch และความยาวลำดับที่แตกต่างกันปนกัน Paged Attention ยืมแนวคิดการออกแบบของหน่วยความจำเสมือน โดยแบ่ง KV Cache เป็น block ขนาดคงที่ และบันทึก block เฉพาะที่แต่ละลำดับใช้ผ่าน block_tables

อินพุตของมันมักจะประกอบด้วยองค์ประกอบต่อไปนี้:

  • query: query ของ token ปัจจุบัน
  • key_cache / value_cache: แคชที่จัดระเบียบตาม block
  • block_tables: ความสัมพันธ์การแมประหว่างลำดับกับ block ทางกายภาพ
  • seq_lens: ความยาวปัจจุบันของแต่ละลำดับ
  • block_size, max_seq_len, scale และข้อมูลเมตาอื่นๆ

ในฐานะบรรณาธิการบริหารและผู้เชี่ยวชาญด้านการเขียนใหม่ขั้นสูง ฉันได้ดำเนินการเขียนใหม่เชิงลึกและลดความซ้ำซ้อนของบทความทางเทคนิคนี้ตามข้อกำหนดทั้งหมดของคุณอย่างเคร่งครัดแล้ว ต่อไปนี้เป็นเวอร์ชันที่ผ่านการประมวลผลแล้ว

5.2 เวลาคอมไพล์ การสร้างอินสแตนซ์เทมเพลต และรูปทรงที่จำกัด

การใช้งาน CUDA ของ Paged Attention อาศัยเทมเพลตและมาโครนิยามเป็นอย่างมาก เพื่อสร้างอินสแตนซ์ฟังก์ชันเคอร์เนลที่สอดคล้องกันสำหรับชุดค่าผสม head_size และ block_size ที่เฉพาะเจาะจง

ความคิดเห็นในโค้ดระบุอย่างชัดเจนว่า เพื่อลดระยะเวลาการคอมไพล์ จึงคอมไพล์เฉพาะ head size ที่ใช้บ่อยเท่านั้น ในขณะเดียวกัน block size ก็ถูกจำกัดไว้เพียงไม่กี่ชุดค่าผสม เช่น 8, 16, 32

// ที่มา: paged-attention/paged-attention/attention/paged_attention_v1.cu

switch (head_size) {

case 32:  LAUNCH_PAGED_ATTENTION_V1(32);  break;

case 64:  LAUNCH_PAGED_ATTENTION_V1(64);  break;

case 128: LAUNCH_PAGED_ATTENTION_V1(128); break;

case 256: LAUNCH_PAGED_ATTENTION_V1(256); break;

default:

TORCH_CHECK(false, "Unsupported head size: ", head_size);

}


#define CALL_V1_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE)
switch (block_size) {
case 8:  CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE); break;
case 16: CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE); break;
case 32: CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE); break;




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


    

        
☕ สนับสนุนค่ากาแฟทีมงาน

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

        

            PromptPay QR
        

        

            SCAN TO PAY WITH ANY BANK
        

    
                                                        
本文来自网络搜集，不代表คลื่นสร้างอนาคต立场，如有侵权，联系删除。转载请注明出处：https://www.itsolotime.com/th/archives/35122
Like (0)
0 0
Previous 15 hours ago
Next 15 hours ago

相关推荐