CuTeDSL في Perplexity

في Perplexity، نستخدم محرك التقديم المحسّن للوقت التشغيلي الداخلي (ROSE) لخدمة النماذج التي تتراوح من التضمينات إلى نماذج اللغات الكبيرة ذات التريليون معلمة. يقع محركنا خلف جميع واجهات برمجة التطبيقات، بما في ذلك Sonar وSearch والتضمينات، وقد تم بناؤه تمكيناً للباحثين من استكشاف النماذج والبنى الأكثر ملائمة، مع نشر المرشحين الأفضل بسرعة في بيئة الإنتاج. ولتحقيق ذلك، نحتاج إلى أن نكون قادرين على تنفيذ ونشر وتحسين الاستدلال للبنى الحديثة بوتيرة سريعة. في منشور المدونة هذا، نستكشف كيف تتيح لنا نواة CuTeDSL المدمجة في محرك الاستدلال الخاص بنا بناء نواة وحدة معالجة الرسوميات المطلوبة لجلب النماذج إلى أداء الذروة على وحدات معالجة الرسوميات من NVIDIA.

يتم استضافة النماذج المخصصة، التي بنتها فرق مختلفة في Perplexity، محلياً على وحدات معالجة الرسوميات NVIDIA Hopper وBlackwell باستخدام محرك الاستدلال ROSE الذي بني بواسطة فريق الاستدلال بالذكاء الاصطناعي. تم بناء ROSE في الأصل لخدمة نماذج Llama المخصصة بشكل أفضل لكل من فك تشفير اللغة والتصنيف، مع واجهة متوافقة مع خادم استدلال NVIDIA Triton. ب مرور الوقت، تطور ليصبح محرك استدلال كاملاً يخدم نماذج لغوية كبير ه جوهرية، بالإضافة إلى نماذج الترتيب والتصنيف والتقييم وغيرها من النماذج القائمة على المحولات لمجموعة واسعة من الفرق لتشغيل منتجات Perplexity. يتمثل هدف تصميم ROSE في البساطة: ب مرور الوقت، يتم استثمار الجهد للوصول إلى أداء الذروة، ولكن الأولوية الرئيسية هي التكيف مع البنى وحالات الاستخدام المتطورة باستمرار لنماذجنا.

مساحة تهيئة النواة عبر أنواع البيانات، وأبعاد الرؤوس، والثوابت وقت الترجمة، وإعدادات التشغيل.

على مستوى عالٍ، يتكون ROSE من محرك يكيّف النموذج مع واجهة يمكن للعملاء التفاعل معها. المحرك ا ت قابلة لإعادة الاستخدام بدرجة كبيرة، حيث تتعامل مع جدولة الطلبات وتجميعها، واسترداد الرموز المفكوكة أو التضمينات أو النتائج من النماذج الأساسية. يتم التعامل مع تهيئة الأجهزة، وقنوات الاتصال بين العمليات وبين العقد، وتحميل الأوزان بواسطة هذه المحركات. بالنسبة لنماذج اللغات الكبيرة، يكون المحرك مسؤولاً عن التجميع والتجزئة والخذرف. وهو مسؤول أيضاً عن تخصيص تخزين KV، لكل من الانتباه الكامل والخطّي، بالإضافة إلى مطابقة البادئة. محرك ا ت التضمينات أبسط، لأنها تجري التجميع عبر الإنترنت فقط. يتم تحسين التنفيذ بحيث يُجهّز المحرك الجولة التالية من العمل بينما يعمل النموذج على المعجل بالتوازي.

يعرض ROSE مجموعة غنية من الطبقات المخصصة النمذجة، التفاف نواة وحدة معالجة الرسوميات المبنية باستخدام CuTeDSL وTriton وCUDA وCUTLASS وcuBLAS. بالنسبة لعمليات ضرب المصفوفات وتطبيقات الانتباه، نعتمد بشكل أساسي على نواة NVIDIA من CUTLASS وcuBLAS. يمثل باقي خط الانتروبيا، بما في ذلك التضمين، والتطبيع، وMoE، ونواة التنشيط، جزءاً من مكتبة النواة في ROSE. لدعم الأجهزة التي نشغل الاستدلال عليها بشكل أفضل، تم تجهيز ROSE ببدائيات الاتصال الخاصة به بين الأجهزة لتوجيه وإرسال وتجميع MoE. في حين أن النواة الأخرى بسيطة نسبياً، يجب تخصصها بدقة لتحقيق أداء الذروة على جميع الأبعاد الخفية والأجهزة المختلفة التي تعمل عليها. بالإضافة إلى ذلك، نعتمد على التخصص في وقت الترجمة لاستيعاب جميع مقاييس وتعديلات النماذج المختلفة (انحياز الوزن مقابل عدم وجود انحياز الوزن لتطبيع RMS، وظائف التنشيط، وما إلى ذلك). ونتيجة لذلك، قررنا استخدام DSLs لتنفيذها: لقد بنينا ROSE في الأصل باستخدام Triton. ومع ذلك، فقد اعتمدنا منذ ذلك الحين CuTeDSL بيئة أساسية برمجة وحدات معالجة الرسوميات.

من منظور رياضي، عادةً ما تكون نواة الاستدلال بسيطة. ونتيجة لذلك، عند تطوير النواة، لا تكمن التعقيدية عموماً في التعبير عن الحساب المراد إجراؤه، بل في كيفية قيام الأجهزة بذلك بالضبط. لبناء أفضل النواة، نطلب لغة توفر التجريدات الصحيحة فوق الأجهزة الأساسية. تتيح لنا اللغة الصحيحة تحديد العمليات عالية المستوى لإجراء الحساب، مع التخصص والتحسين بقوة عبر جميع النقاط في مساحة تكوين واسعة.

CuTeDSL هي لغة DSL تعتمد على Python لوحدات معالجة الرسوميات من NVIDIA مبنية أعلى جير جبر تخطيط CuTe وMLIR. يتم تجميدها في الوقت المناسب إلى PTX مُحسَّن للغاية، مع الاحتفاظ بالتحكم في بدائيات الأجهزة منخفضة المستوى. تهدف إلى توفير نفس الميزات والأداء لـ CUTLASS C++، دون أوقات الترجمة الباهظة التي تعيق سرعة التطوير. من بين اللغات التي استعرضناها أو عملنا معها (بما في ذلك CUDA وcuTile وTVM IR وTriton)، نعتبرها الأنسب لتطوير النواة الخاصة بنا.

بشكل عام، على الرغم من أنها ليست عالية المستوى والموجزة مثل Triton أو cuTile، وجدنا أن جبر تخطيط CuTe وتركيباتها (المعرضة عبر cute.Tensor) كانت بدائل معقولة. بالإضافة إلى ذلك، توفر CuTeDSL وصولاً غير مقيد إلى بدائيات الأجهزة، مما يوفر لنا مستوى التحكم الذي نحتاجه لتحقيق أداء الذروة في جميع حالات الحواف. باستخدام CUDA، كانت فرص التخصص في وقت الترجمة محدودة، حيث أضاف كل معامل إضافي تكلفة غير معقولة لأوقات الترجمة أو التعقيدية لنظام البناء. في المقابل، مع DSL المترجم في الوقت المناسب، يمكننا التخصص بشكل أكثر عدوانية دون التضحية بوتيرة التكرار. على عكس TVM IR، الذي استخدم CUDA خطوة وسيطة، وجدنا أن خط الانتروبيا المستند إلى MLIR الذي انخفض إلى PTX كان أكثر قوة وأسهل في التصحيح من الخط الذي أصدر CUDA خطوة وسيطة. بالإضافة إلى ذلك، مع استخدام العمليات المخصصة لالتفاف تعليمات PTX، وجدنا أن بإمكاننا الاستفادة الكاملة من جميع ميزات الأجهزة الأساسية.

يُقضى الجزء الأكبر من تكلفة ترجمة نواة CUDA عادةً في مرحلة توسيع قالب النواة، حيث يمكن إكمال التحسين والترجمة من مرحلة PTX فصاعداً في غضون ثوانٍ. مع وجود العديد من معاملات وقت الترجمة، يمكن أن تنفجر هذه المرحلة من الترجمة إلى عدة دقائق لنواة غير تافهة. يتزايد هذا الاختناق بحقيقة أن فرص الترجمة المتوازية تنشأ في وقت متأخر فقط، مستويات IR وPTX، مع التعامل مع قالب C++ بشكل تسلسل. وفي حين توجد حيل لمعالجة ذلك، عبر نقل إنشاءات القوالب إلى وحدات ترجمة منفصلة، فإنها تتطلب تعديلات كبيرة على أنظمة البناء واستخدام وحدات ماكرو ثقيلة. في المقابل، لا تتطلب نواة CuTeDSL التي يتم ترجمتها في الوقت المناسب (JIT) قواعد بناء معقدة ولا تعيق أوقات الترجمة، حيث يمكن تركيز دورة التطوير على نقطة واحدة في المساحة التكوينية، بينما يمكن لاختبارات الوحدة الدقيقة لاحقاً تغطية مساحة التصميم بأكملها للنواة.

على الرغم من أن ترجمة JIT تنقل التكاليف من وقت التطوير إلى وقت التشغيل، إلا أن التكلفة العامة أكثر من مقبولة. عندما يتم نشر النموذج، فإنه ينشط فقط عدد ا قليلاً من النقاط من مساحة التكوين بأكملها، مما يتسبب في تكاليف عامة على مقياس الثواني أثناء بدء تشغيل/إحماء محرك الاستدلال. في المقابل، يمكن أن تمتد الترجمة المسبقة لجميع تكوينات النواة لعدة دقائق. بالإضافة إلى ذلك، يمكن تخفيف التكاليف العامة بشكل أكبر من خلال التخزين المؤقت للأصول المترجمة بـ JIT.

تتزايد تكلفة الترجمة الكبيرة لـ CUDA بسبب صعوبة تتبع أصل الأخطاء عند استخدام الإرسال والتخصص القائم على القالب، كما هو الحال مع CuTe وCUTLASS. مع DSL، يمكن الإشارة إلى الأخطاء سطر الكود الذي أنشأها، مما يسهل تصحيح الأخطاء بشكل كبير. بالإضافة إلى ذلك، من الأبسط بكثير تعيين الكود المُنشأ إلى عبارة DSL التي نشأت عنها عبر مراحل MLIR الوسيطة أو الاعتماد على معلومات تصحيح الأخطاء المرفقة بـ PTX المُنشأ. بفضل هذه الجوانب من DSL، وجدنا أنفسنا أكثر إنتاجية بكثير عند تطوير النواة.

في حين تتشارك النواة في العديد من أوجه التشخي عبر النماذج، يجب أن تظل مجهزة بمجموعة من مقاييس التكوين التي يمكن أن يكون لها تأثير على الأداء. تشمل المعاملات ذات الصلة:

البعد الخفي / بعد الرأس / عدد رؤوس الانتباه

أنواع بيانات الإدخال / الإخراج

مخطط التكميم

ثوابت وقت الترجمة (الانحيازات، وظائف التنشيط، إلخ)

بنية ROSE، بما في ذلك المحرك، والنموذج، ومكتبة النواة، والخدمات.

يعد التخصص في وقت الترجمة مفيداً لكل من الأداء وتصحيح الأخطاء. سيما في المعالج الخطي، كما هو الحال مع وحدات معالجة الرسوميات، فإن الفرع غير الموجود أفضل من الفرع المأخوذ دائماً أو غير المأخوذ أبداً. بالإضافة إلى ذلك، تؤدي ترجمة النواة لمجموعة معينة من مقاييس التكوين إلى تعليمات PTX/SASS أقل، والتي تكون أسهل في التحليل والتصحيح عند استخدام cuda-gdb أثناء التحقيق.

بالنسبة لحجم دفعة وتكوين معينين، يمكن تقدير حركة المرور المتوقعة إلى DRAM وعدد العمليات الحسابية. بعد ذلك، يمكن اتخاذ قرارات مثلى حالة بحالة لتخصيص الشبكة الأمثل والمزامنة واستخدام ميزات معينة للأجهزة. تتيح لنا CuTeDSL تحديد وإعادة استخدام المكونات للتخصص العميق دون تكبد تكلفة عامة غير مقبولة وقت الترجمة.

عموماً، يجب تحسين أداء نواة الاستدلال لحالتي استخدام منفصلتين:

التعبئة المسبقة: حجم دفعة صغير (قطع قليلة، تسلسلات)، العديد من الرموز لكل تسلسل. يجب زيادة الإنتاجية للحد الأقصى لكل من التضمينات ورموز الإدخال لنماذج اللغات الكبيرة. زمن الوصول في حدود عشرات إلى مئات المللي ثانية.

فك الترميز: أحجام دفعات صغيرة إلى كبيرة، من رمز واحد (بدون فك ترميز استباقي) إلى بضعة رموز لكل تسلسل (التخمين من المسودة إلى الهدف، التنبؤ متعدد الرموز). زمن الوصول قريب من زمن الوصول عند إطلاق النواة.

بالنسبة لـ GEMM والانتباه، تشمل مرحلة التعبئة المسبقة عدد ا كبيراً من الرموز ويتقيد الأداء بإنتاجية الحساب للجهاز. يتضمن فك التشكيل نفس العملية على أحجام دفعات صغيرة، مع هيمنة أوقات التنفيذ تكلفة تحميل التنشيط والأوزان من DRAM في المعالجات المتعددة للبث ونوى التنسيق، مقيدة بعرض النطاق الترددي للذاكرة للجهاز بدلاً من ذلك. على الرغم من أن حجم الدفعة أو العدد الإجمالي للرموز في الدفعة مؤشر واضح نسبياً، فإن عتبة التحول بين التعبئة المسبقة وفك التشكيل أكثر دقة لأن عرض النطاق الترددي للذاكرة يتحدد أيضاً البعد الخفي للنموذج ونوع بيانات التنشيط. ونتيجة لذلك، يتم تقسيم النواة إلى تنفيذي التعبئة المسبقة وفك التشكيل إذا كانت الاختلافات الخوارزمية أو تعيينات الشبكة/الكتلة تفضل أياً منهما.

المعامل الذي يختلف عالمياً بين نواة التعبئة المسبقة وفك التشكيل هو عدد الوحدات الملتوية في كتلة (عدد السلاسل، مقسوماً على حجم الالتواء البالغ 32 سلسلة لكل التفاف). من المستحسن جعل هذا المعامل ثابتاً وقت الترجمة، لأنه يمكن أن يساعد في تخصيص السجلات. في CUDA، يتم ذلك عادةً بتحديد عدد السلاسل عبر __launch_bounds__ وتوفير عدد الوحدات الملتوية معامل قالب:

cpp

template <size_t NUM_WARPS>
__global__ void 
__launch_bounds__(NUM_WARPS * WARP_SIZE, 1)
some_kernel(...) {

}

عدد السجلات المتاحة لجميع السلاسل في كتلة هو قيد على الأجهزة، 64K على Hopper وBlackwell. إذا كان عدد السلاسل معروفاً مسبقاً، يعرف المترجم بالضبط عدد السجلات التي يمكنه تعيينها لكل سلسلة ويمكنه نجاح ترجمة نواة محسنة بشكل أكثر عدوانية. ومع ذلك، يأتي هذا بتكلفة في وقت الترجمة، حيث يتعين الآن إعادة ترجمة النواة لأحجام التفاف مختلفة. مع CuTeDSL، لم تعد هذه مشكلة، لأننا لا نحتاج إلى دفع تكلفة الترجمة لمساحة التكوين بأكملها مقدماً. تتوقع DSL أيضاً أن يكون حجم الكتلة تعبيراً ثابتاً، مما يؤدي ضمناً إلى تنفيذ التحسين الذي تعنيه __latch_bounds__ في CUDA.

إلى جانب جعل عدد الوحدات الملتوية ثابتاً، من المرغوب فيه تباينه بين التعبئة المسبقة وفك التشكيل. للاستفادة الكاملة من عرض النطاق الترددي للذاكرة لوحدة معالجة الرسوميات، لا يكفي دائماً تعيين عدد كبير من الوحدات الملتوية لكل كتلة، حيث لا يمكن تحقيق عرض النطاق الترددي الكامل إلا باستخدام أكبر عدد ممكن من المعالجات المتعددة للبث المتاحة (SMs). ونتيجة لذلك، لعدد معين من الرموز وأنواع البيانات المعروفة والأبعاد الخفية، نختار أولاً بُعد متجه مثالي لكل سلسلة (8 لـ bfloat16، 4 لـ float32) لضمان استخدام الكود المُنشأ لتعليمات التحميل/التخزين 128 بت (256 بت على Blackwell). بمعرفة حجم المتجه، يمكننا حساب عدد الوحدات الملتوية المطلوبة بواسطة النواة. نحاول أولاً زيادة عدد SMs للحد الأقصى: إذا كان هناك عدد أقل من الوحدات الملتوية من SMs، يتم اختيار التفاف واحد لكل كتلة. هذا هو عادةً البديل لفك التشكيل للنواة. إذا كان هناك عدد كافٍ من الوحدات الملتوية، فنحن نزيد تدريجياً عدد الوحدات الملتوية إلى 8 أو 16 أو 32، مع قياس معيار للحد الأعلى لم يعد بعده نلاحظ عادةً تحسناً في الأداء.

في حين قام Triton ببعض هذه التحسينات تحت الغطاء، وجدنا أنه يمكننا تحقيق نتائج مماثلة باستخدام جبر تخطيط CuTe بعناية، دون تقييدنا عيوب المترجم عالي المستوى. سيما عندما تشمل الاختزال ات الأفقية دون الالتفاف، وجدنا أن CuTeDSL قدمت لنا التحكم الدقيق المطلوب لإنشاء كود مثالي.

في بعض الحالات، لتوزيع العمل عبر أكبر عدد ممكن من الكتل، يمكن استخدام cooperative_groups::this_grid().sync() لمزامنة جميع الكتل المشاركة. تنطبق قيود: نظراً لعدم إمكانية إلغاء جدولة كتلة من SM بمجرد بدء تنفيذها، لضمان التقدم للأمام، لا يمكن أن يتجاوز عدد الكتل عدد SMs المتاحة على الجهاز، لضمان قدرة جميع الكتل في أي نقطة على تحديث واستطلاع مواقع الذاكرة العامة المشاركة في المزامنة. يتم فرض هذا القيد بواسطة cudaLaunchCooperativeKernel. تبلغ تكلفة حاجز الشبكة حوالي 3 ميكروثانية على Hopper، بالإضافة إلى تكلفة مسح جميع عمليات الذاكرة السابقة عبر حاجز ذاكرة ضمني في بعض الحالات. بشكل عام، هذا مشابه للتكلفة العامة لتقسيم النواة إلى عمليتي إطلاق، مما يؤدي إلى قرار تصميم حساس للأداء. على الرغم من حقيقة أن الإطلاق التابع البرمجي (PDL) يمكن أن يقلل التكاليف العامة للإطلاق من Hopper فصاعداً، بالنسبة لنواة فك التشكيل، فإن تجنب التكلفة العامة يستحق العناء.

إذا تطلبت إحدى أحمال العمل كفالة أكثر من عدد SMs، فيجب تقسيم النواة، بالاعتماد على الحاجز الضمني لإطلاق النواة. هذا هو الحال الأساسي للتعبئة المسبقة. إذا كان عدد الكتل المثالي لجميع أحمال العمل يناسب الشبكة، يُفضل حاجز الشبكة، كما هو الحال مع نواة فك التشكيل التي تعمل على عدد أقل من الرموز.

في CuTeDSL، يمكن اتخاذ هذه القرارات في كود الإطلاق من جانب المضيف ويمكن تركيب النواة وتداخلها مع حواجز الشبكة حسب الضرورة، كما يوضح المثال أدناه. مع هذه الاستراتيجية الهجينة، يمكن تحسين نواة فك التشكيل للوصول إلى الحد الأدنى من زمن الوصول ويمكن تحسين نواة التعبئة المسبقة للإنتاجية القصوى. على عكس الحلول المستندة إلى CUDA، يمكن تحديد هذا المنطق بشكل أكثر إيجازاً وبأقل تكلفة عامة وقت الترجمة.

python

@cute.kernel
def dispatch_kernel_decode(...):
	route_tokens(...)
	grid_barrier(grid_barrier_ptr)
	copy_tokens_to_experts(...)

@cute.kernel
def dispatch_kernel_prefill_route(...):
	route_tokens(...)
	
@cute.kernel
def dispatch_kernel_prefill_copy(...):
	copy_tokens_to_experts(...)
	
@cute.jit
def dispatch(...):
	num_route_blocks = ...
	num_copy_blocks = ...
	
	if num_copy_blocks < NUM_SMS:
		dispatch_kernel_decode(...).launch(
			grid=(min(max(num_route_blocks, num_copy_blocks), NUM_SMS), 1, 1)
		)
	else:
		dispach_kernel_prefill_route(...).launch(
			grid=(num_route_blocks, 1, 1)
		)
		dispatch_kernel_prefill_copy(...).launch(
			grid=(num_copy_blocks, 1, 1)
		)

في حالة نواة MoE all-to-all الخاصة بنا، كما هو موضح بالتفصيل أدناه، ساعدنا توفر مرونة بناء النواة مع أو بدون حواجز شبكية وتحسينها وقت الترجمة في تحقيق مكاسب كبيرة في الأداء عبر CuTeDSL.

كانت النواة الأولى التي بنيناها في CuTeDSL هي نواة تطبيع RMS مخصصة، محسنة للأبعاد الصغيرة. يُطبق تطبيع QK، المستخدم في طبقة الانتباه الذاتي لنماذج Qwen الحديثة، أساس لكل رأس، على 64 أو 128 أو 256 عنصراً. على عكس تطبيع RMS العادي، الذي يعمل على الأبعاد الخفية التي تتجاوز الآلاف، ويتطلب اختزالاً أفقياً عبر كل من الذاكرة المشتركة والسجلات، يحتاج تطبيع QK إلى الاختزال عبر الوحدات الملتوية فقط، مع إمكانية ضغط رؤوس متعددة في التفاف واحد. أصبح هذا حالة استخدام حيث يمكننا التخصص بقوة وقت الترجمة للعثور على أفضل تعيينات للرموز إلى الوحدات الملتوية لجميع أبعاد الرأس وأنواع البيانات.

تطبيع RMS هو عملية واضحة تماماً، على الرغم من أنها تتطلب مجموعاً أفقياً لحساب المعيار المربع للمتجه:

cpp

r = w * (x * ((x * x).sum(dim=-1) / x.shape[-1] + EPS))

نظراً لأن رؤوس الانتباه صغيرة عادةً ويتم دمج عمليات كل سلسلة بشدة، يمكن لالتفافة واحدة حساب المعيار المربع. لزيادة عرض النطاق الترددي للذاكرة للحد الأقصى، يجب أن تعالج كل سلسلة 128 بيتاً من الإدخال، وهو ما يترجم إلى 8 قيم bfloat16. ونتيجة لذلك، تكون نصف التفافة كافية لمعالجة الرأس، ويمكن تحديد المعيار المربع باستخدام اختزال أفقي يتم إجراؤه بواسطة shuffle_sync_bfly. بدءاً من بُعد الرأس ونوع البيانات، يمكننا تحديد عدد الرؤوس لكل التفاف وتحديد المعاملات وعدد جولات الاختزال وقت الترجمة.

يمكننا بعد ذلك بناء نواة مخصصة بعدد قليل نسبياً من أسطر الكود، بالاعتماد على عمليات التحميل/التخزين الموجهة تلقائياً واستخدام عمليات خاصة بالبنية التحتية منخفضة المستوى لاختزال الالتفاف. تحدد النواة الرأس (Q أو K) المراد العمل عليه وتحدد الأوزان المقابلة لتطبيقها. وهي تحمل متجهاً كاملاً من الرأس وتجمع محلياً أولاً. بعد ذلك، بناءً على عدد السلاسل لكل رأس، يتم حساب المعيار المربع. أخيراً، يتم تحديد تطبيع RMS، الموجه تلقائياً، والكتابة به مرة أخرى إلى الذاكرة. يتم طي الثوابت المعاملات الخاصة بالنموذج، مثل WEIGHT_BIAS وEPS.

python

@cute.kernel
def _qk_norm_kernel(
    q: cute.Tensor,
    k: cute.Tensor,
    q_weight: cute.Tensor,
    k_weight: cute.Tensor,
    WEIGHT_BIAS: cutlass.Constexpr[float],
    EPS: cutlass.Constexpr[float],
) -> None:
    # Find the block and thread indices.
    bidx, _, _ = cute.arch.block_idx()
    bdim, _, _ = cute.arch.block_dim()
    tidx, _, _ = cute.arch.thread_idx()

    _, num_q_heads, (threads_per_head, vecs_per_thread, vec_size) = q.shape
    _, num_k_heads, _ = k.shape
    head_dim = threads_per_head * vecs_per_thread * vec_size

    # Determine the thread and the vector to work on.
    vec_id = (bidx * bdim + tidx) // threads_per_head
    thread_id = tidx % threads_per_head

    # Determine the token and head indices.
    token = vec_id // (num_q_heads + num_k_heads)
    head_offset = vec_id % (num_q_heads + num_k_heads)
    is_q = head_offset < num_q_heads

    # Determine which tensor to work with.
    thread_coord = (thread_id, None, None)
    head_coord = (token, head_offset if is_q else head_offset - num_q_heads, thread_coord)

    # Load the head into registers.
    x_ref = q[head_coord] if is_q else k[head_coord]
    w = (q_weight if is_q else k_weight)[thread_coord].load()

    # Compute the squared sum locally.
    x = x_ref.load()
    x_f32 = x.to(cute.Float32)
    x_square = x_f32 * x_f32
    x_square_sum = x_square.reduce(cute.ReductionOp.ADD, init_val=0.0)

    # Sum up across the warp.
    mask = activemask()
    start = cute.arch.WARP_SIZE // (cute.arch.WARP_SIZE // threads_per_head * 2)
    for offset in [start >> i for i in range(5) if (start >> i) > 0]:
        x_square_sum += cute.arch.shuffle_sync_bfly(
            x_square_sum,
            offset=offset,
            mask=mask,
        )

    # RMS norm.
    rms = cute.math.rsqrt(x_square_sum / head_dim + EPS)
    x_norm = ((w + WEIGHT_BIAS) * (x * rms)).to(x.element_type)
    x_ref.store(x_norm)

قارناً بنواة تطبيع RMS العادية، يوفر هذا البديل المخصص إنتاجية وزمن وصول أفضل بـ 2-3 أضعاف لكل من التعبئة المسبقة وفك التشكيل.

يقوم إرسال وتجميع MoE بخلط رموز الإدخال إلى الخبراء الذين تم توجيههم إليهم، وإجراء تراكم مرجح بعد أن تجري نواة MoE GEMM عمليات ضرب المصفوفات ذات الصلة. الأداء محل اهتمام خاص، حيث تعتمد منتجات مختلفة على نماذج MoE بأحجام مختلفة. بناءً على أنواع النشر، نميز بين نواة مختلفة للتحسين:

TP=1: نواة محلية لوحدة معالجة رسوميات واحدة. يخلط الإرسال الرموز إلى الخبراء المعنيين ويقوم التجميع بتراكمها. تُستخدم بشكل أساسي لنشر Qwen3 30B على H100.

DP=1, TP ≠ 1: نواة متوازية التنسيق، أحادية العقدة. يتم نسخ الإدخال، ويختار الإرسال الرموز للرتبة الحالية، ويقوم التجميع بتراكمها عبر جميع الرتب. التراكم عبر NVLink أو NVLS. تُستخدم لعمليات التعبئة المسبقة الكبيرة.

DP ≠ 1, NVLink: نواة متوازية البيانات مع توازي التنسيق الاختياري. ترسل كل رتبة رموزاً إلى كل رتبة أخرى تستضيف خبراء محددين. يتم تجميع معلومات التوجيه، ونقل الرموز عبر NVIDIA NVLink. تُستخدم لنشر H200 أو GB200 المتوازي للبيانات.

DP ≠ 1, InfiniBand: نواة متوازية البيانات ومتوازية التنسيق من نظير إلى نظير تعمل عبر كل من NVLink وNVIDIA Quantum InfiniBand. تُستخدم لنشر Hopper متعدد العقد.

تحسب النواة المحلية عدد الرموز المعينة لكل خبير بناءً على جدول التوجيه وتعين فهرساً فريداً لكل رمز داخل قائمة الرموز المنتمية لكل خبير. بعد خطوة التوجيه، يتم نسخ الرموز إلى مخازن مؤقتة مبطنة اختيارياً ليتم تمريرها إلى نواة MoE GEMM. باستخدام مواقع الرموز، تجري نواة التجميع بعد ذلك تراكماً في السجلات، وتكتب الرموز النهائية في مخزن مؤقت للإخراج.

تختلف مرحلة الإرسال بين التعبئة المسبقة وفك التشكيل: بالنسبة للرموز القليلة، هناك القليل من معلومات التوجيه لقراءتها، لذلك نقوم بمعالجتها في كتلة واحدة والإشارة إلى باقي الشبكة، مع نسخ كتل فردية لقطع من الرموز. يتم تحديد تعيين الرموز بزيادة العدادات ذرّياً في الذاكرة المشتركة. نظراً لأن التعبئة المسبقة يجب أن تعالج المزيد من معلومات التوجيه، يتم تقسيم التوجيه عبر كتل متعددة في نواة منفصلة. يتم تعيين الفهارس محلياً أولاً باستخدام العدادات في الذاكرة المشتركة ثم يتم تجميعها عبر جميع الكتل في الذاكرة العامة لتحديد موقع كل رمز. يتم تثبيت شبكة هذه النواة على عدد SMs لكي تتمكن من استخدام حواجز على مستوى الشبكة. تجري نواة لاحقة النسخ، ويتم إطلاقها عبر عدد أكبر بكثير من الكتل والوحدات الملتوية للاستفادة الكاملة من عرض النطاق الترددي المتاح للذاكرة.

نواة التجميع بسيطة نسبياً، حيث تجري تراكماً موجهاً، بالاعتماد على معلومات التوجيه والإزاحات المحسوبة بواسطة الإرسال. أثبتت CuTeDSL قيمتها هنا، حيث يمكن التعبير عن الحساب من خلال موترات CuTe البسيطة التي تترجم إلى كود فعال لجميع أنواع البيانات والأبعاد الخفية.

في نشر متوازي التنسيق، يتم نسخ معلومات التوجيه ورموز الإدخال. يمكن إعادة استخدام نواة الإرسال المحلي بالكامل مع تغيير طفيف لتوجيه واختيار الرموز المعينة حصرياً للخبراء الموجودين على الرتبة الحالية. مع بعض المعاملات والتخصص وقت الترجمة، يتم إعادة استخدام نواة الإرسال المحلي.

نظراً لأن عدد الرتب عادةً أقل من عدد الخبراء الذين يتم توجيه الرمز إليهم، يمكن افتراض في المتوسط أن كل رمز لديه نسخة واحدة على الأقل من نسخه الموجهة إلى كل رتبة. في هذه الحالة، يمكن إعادة استخدام نواة التجميع المحلية لتجميع النسخ المحلية للرمز، يليه تقليل شامل (all-reduce) للجمع عبر الرتب المشاركة. في حالة نواة التعبئة المسبقة، نطلق نواة منفصلة، ومع ذلك يدمج فك التشكيل كلتا العمليتين في إطلاق واحد.

بالنسبة للتقليل الشامل (all-reduce)، نستخدم NVLink العادي إذا شاركت رتبتان فقط، مع التبديل إلى NVLS عند مشاركة 4 رتب أو أكثر بزيادة قدرها 10-20% في الإنتاجية. بفضل DSL، يمكننا مشاركة معظم منطق الاختزال وتبادلات المخزن المؤقت للنظير بين MoE ونواة التقليل الشامل العادية. هذا موقف يكون فيه التحكم الدقيق في الأجهزة مفيداً، حيث تتطلب عمليات DSL عبر multimem.ld_reduce وmultimem.st لـ NVLS. يتبع تنفيذ الاختزال نفسه إلى حد كبير مثال all_reduce من اختبارات NCCL.

تعتمد نواة متوازي البيانات على NVLink لتجميع الرموز لأجهزة متعددة. في مرحلة الإرسال، قبل تبادل الرموز، توجّه النظائر محلياً أولاً لتحديد الإزاحة المستهدفة على كل رتبة لكل رمز ترسله. بعد ذلك، تبدأ في نسخ الرموز إلى المخازن المؤقتة للاستلام المخصصة لكل مرسل لتجنب التعارضات. تُستخدم كتابات تحرير-تخزين لكل رمز للإشارة إلى الاكتمال. من جانب المستلم، يتم تجميع معلومات التوجيه لتكون قادرة على ضغط جميع الرموز في مخزن مؤقت واحد لـ MoE GEMM. يستطلع المستعمرون العلامة التي كتبها النظير بعملية جلب-تحميل (acquire-load) لانتظار اكتمال عمليات الكتابة.

من جانب التجميع، تُنسخ الرموز عبر NVLink إلى رتب المصدر، مع علامة لكل رمز تشير إلى الاستلام. نواة التجميع مشابهة للنسخ المحلية وأحادية العقدة، مع اختلافات طفيفة في العنونة وبعض منطق المزامنة الإضافي.

لقد بنينا مكتبة من غلافات PTX لعرض عمليات جلب وتحميل التحرير والاستحواذ ذات النطاق على مستوى النظام والجهاز لتنفيذ هذه النواة. من خلال تحسين دقة علامات المزامنة وتخصيص المزيد من الكتل لإرسال واستلام الرموز عبر NVLink، حققنا تحسينات بمقدار ضعفين مقارنة بنواة CUDA السابقة.

حققت نواة النظير إلى النظير الخاصة بنا سابقاً أداءً متطوراً على وحدات معالجة الرسوميات Hopper المقترنة بمحولات ConnectX-7 InfiniBand. لقد قمنا منذ ذلك الحين بنقل نواة التجميع إلى CuTeDSL لتحقيق مكاسب إضافية في الأداء. نظراً لأن إصدارات CUDA كانت تعتمد بالفعل بشدة على التخصص وقت الترجمة، فقد ثبت أنها اختناقات وقت الترجمة في مستودعنا، مما أجبرنا على حذف بعض المعاملات. من خلال الانتقال إلى DSL المترجم بـ JIT، تخصصنا لمزيد من المعاملات وعدلنا تخصيصات الشبكة، مما أدخل اختلافات إضافية بين التعبئة المسبقة وفك التشكيل لزيادة زمن الوصول بحوالي 3%.

من منظور نواة ROSE، ينقسم دعم Blackwell إلى فئتين: نواة محايدة للبنية ونواة خاصة للبنية. تتطلب النواة الخاصة للبنية، التي تعتمد على نواة التنسيق، عملاً صريحاً للتكيف. ومع ذلك، بفضل استخدام DSL، تتكيف النواة المحايدة للبنية بسلاسة.

تحتوي وحدة معالجة الرسوميات Blackwell على ما يصل إلى 160 SMs، على عكس 132 في Hopper. بالنسبة لمعظم النواة، هذا ثابت وقت الترجمة والمنطق في كود المُطلق من جانب المضيف يأخذه في الاعتبار لتحديد النواة المثلى (التعبئة المسبقة مقابل فك التشكيل) للتشغيل. بالإضافة إلى ذلك، تقدم Blackwell أيضاً تعليمات تحميل 256 بت، والتي يمكن إرسارها بسهولة زيادة عرض المتجه (البعد الأخير لموتر CuTe) لتحسين استغلال عرض النطاق الترددي للذاكرة.

تقدم Blackwell دعماً في الأجهزة لعدد كبير من مخططات التكميم (MXFP8 وMXFP4 وNVFP4). تم تنفيذ دعم هذه المخططات أساساً عن طريق تعديل النواة الحالية لتحجيم كتلة FP8. يُعد دعم JIT أمحاً حاسماً هنا، حيث يمكن دعم أوضاع جديدة دون انفجار مساحة التكون وقت الترجمة. في الماضي، وجدنا أن Triton غير ملائم في التعبير عن الاختزال ات الأفقية عبر عناصر قليلة، فوق نصف أو ربع التفافة. كان التحكم الدقيق في الاختزال ات أمراً بالغ الأهمية لتحقيق الأداء الأمثل عند حساب قيمة حد أقصى مطلقة عبر 16 أو 32 عنصراً لـ NVFP4 وMXFP8/MXFP4، على التوالي.

باعتبارها لغة جديدة نسبياً، تلقت نماذج اللغات الكبيرة تعرضاً ضئيلاً لكود CuTeDSL أثناء التدريب المسبق. بالإضافة إلى ذلك، كان نمط كودنا الحالي، مع النواة المحسنة للتنشيط والتطبيع، مختلفاً قليلاً عن أمثلة CUTLASS أو Quack الحالية، والتي تم تحسينها بشدة لتعليمات ضرب المصفوفات الخاصة بـ TMA والبنية التحتية. كان الوصول إلى نقطة يمكن لنماذج اللغات الكبيرة المساعدة في كتابة النواة عملية تدريجية.

قمنا يدوياً ببناء مجموعة أولية من النواة ووظائف الأداة المساعدة، باتخاذ أنماط متسقة، قبل أن نتمكن من الاعتماد على الوكلاء لتحويل النواة. مع تقدم المجموعة، يمكننا إجراء ترجمات تقريبية من CUDA إلى CuTeDSL، يعقبها إعادة صياغة يدوية كبيرة. على وجه الخصوص، وجدنا أن النماذج عانت مع تحويلات المؤشر إلى الموتر والتوجيه، وعلقت في محاولة إصلاح أخطاء النوع عند محاولة جعل الاختبار يترجم. لقد أظهرت فهماً ضئيلاً للطبيعة القائمة على الشبكة للنواة، حيث خلطت بين مزامنة السلسلة والشبكة. والأكثر إحباطاً، كان هناك إغراء هلوسة ترجمات غير موجودة لمكتبات خاصة بـ CUDA بدلاً من البناء مقابل بدائل CuTeDSL الخاصة بنا. ب مرور الوقت، مع نمو مجموعة نواة CuTeDSL، وجدنا أن نماذج اللغات الكبيرة يمكنها بناء معظم كود الهيكل وملء الفراغات، مما يتيح لنا التركيز على التصميم عالي المستوى.

مقارنة بـ CUDA، نجد أن CuTeDSL تقدم أداء أفضل بكثير وقت الترجمة، مما يتيح لنا التخصص لمزيد من المعاملات وتحقيق تحسينات في الأداء. وفي حين أن النواة أكثر إطناباً قليلاً بسبب نقص تجريدات معين مقارنة بـ Triton، نجد أن التعقيدية مبررة حيث يتيح لنا التحكم الدقيق في الأجهزة بناء نواة أسرع.

وقت الكتابة، الفجوة في الأداء بين CUTLASS C++ ونواة CuTeDSL ضئيلة، مما يفتح فرصاً مثيرة لتركيب ودمج نواة GEMM عالية الأداء مع الاتصال بين الأجهزة أو التطبيع أو التنشيط. بالالتزام بـ DSL هذه، سنتمكن من استخراج المزيد من الأداء من أجهزتنا لتحسين الإنتاجية عبر أحمال العمل لدينا، من التضمينات إلى نماذج اللغات الكبيرة.

https://github.com/Dao-AILab/quack

https://github.com/NVIDIA/cutlass

https://github.com/NVIDIA/nccl-tests

https://docs.nvidia.com/cuda/parallel-thread-execution/

BibTeX

@online{perplexity2026CutedslAtPerplexity,
  author = {{Perplexity} and Nandor Licker and Bex Brooker and Jesse Dwyer and Phi Hoang},
  title = {{CuTeDSL at Perplexity}},
  date = {2026-05-06},
  url = {https://www.perplexity.ai/hub/blog/cutedsl-at-perplexity}
}

الاستدلال في Perplexity AI

لماذا CuTeDSL؟

تجربة المطور

مساحة تكوين كبيرة

التعبئة المسبقة وفك تشفير التخصص

مزامنة الشبكة

نواة CuTeDSL

تطبيع QK

إرسال/تجميع خليط الخبراء (Mixture-of-Experts)

محلي

أحادية العقدة

متوازي البيانات

من نظير إلى نظير

دعم Blackwell

إنشاء النواة بمساعدة الذكاء الاصطناعي

الخلاصة والعمل المستقبلي

المراجع

اقتباس