با افزایش تقاضا برای مدلهای زبانی بزرگ (LLM) محلی، توسعهدهندگان از کتابخانههای استاندارد مانند PyTorch و vLLM فراتر رفته تا پایپلاینهای استنتاجی بسیار سفارشی بسازند. در حالی که کتابخانههای استاندارد پیشفرضهای عالی ارائه میدهند، آنها اغلب بهینهسازیهای ریز را به خاطر قابلیت حمل و سهولت استفاده فدا میکنند. برای بیرون کشیدن هر قطره از توان پردازشی از کارتهای گرافیک سطح مصرفکننده یا A100های دیتاسنتر، باید به عمق بهینهسازی هستههای CUDA بپردازید. این پست تکنیکهای حیاتی برای تنظیم هستههای CUDA را که به طور خاص برای معماریهای استنتاج ترنسفورمر طراحی شدهاند، بررسی میکند.
درک گلوگاه سلسله مراتب حافظه
گلوگاه اصلی در استنتاج LLM به ندرت محدود به محاسبات است؛ بلکه محدود به حافظه است. مکانیسمهای توجه و لایههای پیشخور، نیاز به جابجایی حجم عظیمی از داده بین حافظه با پهنای باند بالا (HBM) و حافظه روی تراشه GPU دارند. اولین گام در تنظیم، اطمینان از همپیوستگی حافظه است. وقتی رشتههای مجاور در یک وارپ به آدرسهای حافظه مجاور دسترسی دارند، سختافزار این درخواستها را در یک تراکنش حافظه واحد ترکیب میکند. عدم همترازی ساختارهای داده یا الگوهای دسترسی میتواند عملکرد را تا یک مرتبه بزرگی کاهش دهد.
یک هسته ضرب ماتریسی استاندارد را در نظر بگیرید. اگر شاخصهای رشته به دقت نگاشت نشوند، ممکن است عملیات پخش گرانقیمتی را متحمل شوید. همیشه اطمینان حاصل کنید که ابعاد بلوک رشته شما مضربی از اندازه وارپ (معمولاً ۳۲) باشد و دسترسیهای حافظه جهانی پیوسته باشند.
استراتژیک استفاده از حافظه مشترک
حافظه مشترک، SRAM با سرعت بالا که روی تراشه GPU قرار دارد، تأخیری چندین مرتبه بزرگی کمتر از حافظه جهانی ارائه میدهد. برای هستههای سفارشی LLM، مانند آنهایی که Flash Attention بهینهشده یا RoPE (جایگذاری موقعیتی چرخشی) را پیادهسازی میکنند، استفاده از حافظه مشترک برای عملکرد غیرقابل مذاکره است.
با این حال، حافظه مشترک یک منبع کمیاب است. تخصیص بیش از حد منجر به ریختن رجیستر میشود که عملکرد را نابود میکند. کلید کار، کاشیبندی (Tiling) است. شما باید دادههای خود را طوری کاشیبندی کنید که در محدودیتهای حافظه مشترک جا شوند و همزمان موازیسازی را به حداکثر برسانید. در زیر نمونه سادهای از نحوه اعلام و استفاده از حافظه مشترک برای بارگذاری ماتریسی مبتنی بر کاشی آورده شده است:
// Inside the CUDA kernel
extern __shared__ float sdata[];
// Load data into shared memory
unsigned int tid = threadIdx.x;
unsigned int row = blockIdx.y * blockDim.y + threadIdx.y;
unsigned int col = blockIdx.x * blockDim.x + threadIdx.x;
if (row < M && col < N) {
sdata[tid] = globalMatrix[row * N + col];
}
__syncthreads(); // Ensure all threads have loaded their data
// Perform computation using shared memory
float sum = 0.0f;
for (int i = 0; i < TILE_SIZE; ++i) {
sum += sdata[tid] * otherMatrix[i * N + col];
}
در این قطعه کد، __syncthreads() حیاتی است. این تابع به عنوان یک مانع عمل میکند و اطمینان حاصل میکند که تمام رشتههای بلوک قبل از اینکه هر رشته شروع به خواندن از آن کند، بارگذاری دادهها را در sdata به پایان رساندهاند. بدون این، شرایط مسابقه محاسبات شما را مخدوش خواهد کرد.
اشغال و فشار رجیستر
اشغال به نسبت وارپهای فعال در هر مولتیپردازنده به حداکثر تعداد ممکن وارپها اشاره دارد. اشغال بالا با اجازه دادن به برنامهریز برای سوییچ کردن به وارپ دیگر در حالی که یک وارپ در انتظار حافظه است، به پنهان کردن تأخیر حافظه کمک میکند. با این حال، افزایش اشغال یک عملیات تعادلی است. هر رجیستر استفاده شده در هر رشته، تعداد کل وارپهایی که میتوانند در یک مولتیپردازنده جا شوند را کاهش میدهد.
برای تشخیص فشار رجیستر، از nvprof یا ncu (NVIDIA Nsight Compute) استفاده کنید. اگر هسته شما به دلیل "محدودیتهای رجیستر" اشغال پایینی گزارش میدهد، باید کد خود را بازسازی کنید. استراتژیهای رایج شامل موارد زیر است:
- ارتقای متغیرهای محلی به حافظه
__shared__. - کاهش اندازه آرایهها و حلقههایی که به فضای پشته قابل توجهی نیاز دارند.
- استفاده از اشارهگرهای
__restrict__برای کمک به کامپایلر در بهینهسازی محاسبات اشارهگر.
استفاده از هستههای تنسور برای دقت ترکیبی
کارتهای گرافیک NVIDIA مدرن دارای هستههای تنسور هستند، واحدهای سختافزاری تخصصی که برای ضربهای ماتریسی با دقت ترکیبی (FP16، BF16، INT8) طراحی شدهاند. برای استنتاج LLM که اغلب نسبت به از دست دادن دقت تحملپذیر است، تبدیل هستههای matmul خود برای استفاده از هستههای تنسور از طریق CUTLASS یا دستورات __mma دستی میتواند سرعتبخشی ۴ تا ۱۰ برابری نسبت به هستههای CUDA استاندارد FP32 یا FP16 به همراه داشته باشد.
هنگام نوشتن هستههای سفارشی، اطمینان حاصل کنید که چیدمان داده شما NHWC (یا مناسب برای هستههای تنسور) باشد، نه NCHW، و ابعاد ماتریس خود را به مضربهای ۸ یا ۱۶ بسته به معماری خاص (Volta، Ampere یا Hopper) همتراز کنید.
نتیجهگیری
تنظیم هستههای CUDA برای استنتاج LLM محلی کار افراد ضعیفالقلب نیست. این کار نیازمند درک عمیقی از معماری GPU، سلسله مراتب حافظه و بهینهسازیهای کامپایلر است. با این حال، پاداشها چشمگیر هستند. با تسلط بر همپیوستگی حافظه، کاشیبندی حافظه مشترک، مدیریت اشغال و استفاده از هستههای تنسور، میتوانید پایپلاینهای استنتاجی بسازید که از چارچوبهای همهمنظوره پیشی بگیرند. با پروفایل کردن گلوگاه خود شروع کنید، این تکنیکها را به صورت تکراری اعمال کنید و همیشه تأثیر آنها را با ابزارهایی مانند Nsight Compute اندازهگیری کنید. آینده هوش مصنوعی محلی کارآمد در دستان کسانی است که میتوانند زبان GPU را صحبت کنند.